<commit_before>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, error, info, o, Logger};
use tokio::io::Result;
use tokio::net::udp::{RecvHalf, SendHalf};
use tokio::net::UdpSocket;
use tokio::sync::mpsc::{channel, Receiver, Sender};
use tokio::sync::{Mutex, RwLock};

use crate::config::{Config, ConnectionConfig};

type SessionMap = Arc<RwLock<HashMap<(SocketAddr, SocketAddr), Mutex<Session>>>>;

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    // TODO: write tests for this
    /// start the async processing of incoming UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        self.log_config(&config);

        let (mut receive_socket, send_socket) = Server::bind(&config).await?.split();
        // HashMap key is from,destination addresses as a tuple.
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, receive_packets) = channel::<Packet>(1024);

        let log = self.log.clone();
        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, receive_packets).await;
            debug!(log, "Receiver closed");
        });

        loop {
            if let Err(err) = Server::process_receive_socket(
                self.log.clone(),
                config.clone(),
                &mut receive_socket,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
            {
                error!(self.log, "Error processing reeive socket: {}", err);
            }
        }
    }

    /// process_receive_socket takes packets from the local socket and asynchronously
    /// processes them to send them out to endpoints.
    async fn process_receive_socket(
        log: Logger,
        config: Arc<Config>,
        receive_socket: &mut RecvHalf,
        sessions: SessionMap,
        send_packets: Sender<Packet>,
    ) -> Result<()> {
        let mut buf: Vec<u8> = vec![0; 65535];
        let (size, recv_addr) = receive_socket.recv_from(&mut buf).await?;
        tokio::spawn(async move {
            let endpoints = config.get_endpoints();
            let packet = &buf[..size];

            debug!(
                log,
                "Packet Received from: {}, {}",
                recv_addr,
                from_utf8(packet).unwrap()
            );

            for (_, dest) in endpoints.iter() {
                if let Err(err) = Server::ensure_session(
                    &log,
                    sessions.clone(),
                    recv_addr,
                    *dest,
                    send_packets.clone(),
                )
                .await
                {
                    error!(log, "Error ensuring session exists: {}", err);
                }

                let map = sessions.read().await;
                let mut session = map.get(&(recv_addr, *dest)).unwrap().lock().await;
                if let Err(err) = session.send_to(packet).await {
                    error!(log, "Error ensuring sending packet: {}", err)
                }
            }
        });
        return Ok(());
    }

    /// process_receive_packet_channel blocks on receive_packets.recv() channel
    /// and sends each packet on to the Packet.dest
    async fn process_receive_packet_channel(
        log: &Logger,
        mut send_socket: SendHalf,
        mut receive_packets: Receiver<Packet>,
    ) {
        while let Some(packet) = receive_packets.recv().await {
            debug!(
                log,
                "Sending packet back to origin";
                "origin" => packet.dest,
                "contents" => String::from_utf8(packet.contents.clone()).unwrap(),
            );

            if let Err(err) = send_socket
                .send_to(packet.contents.as_slice(), &packet.dest)
                .await
            {
                error!(log, "Error sending packet"; "dest" => packet.dest.to_string(), "err" => err.to_string());
            }
        }
    }

    /// log_config outputs a log of what is configured
    fn log_config(&self, config: &Arc<Config>) {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// ensure_session makes sure there is a value session for the name in the sessions map
    async fn ensure_session(
        log: &Logger,
        sessions: SessionMap,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<()> {
        {
            let map = sessions.read().await;
            if map.contains_key(&(from, dest)) {
                return Ok(());
            }
        }
        let s = Session::new(log, from, dest, sender).await?;
        {
            let mut map = sessions.write().await;
            map.insert((from, dest), Mutex::new(s));
        }
        return Ok(());
    }
}

/// Packet represents a packet that needs to go somewhere
struct Packet {
    dest: SocketAddr,
    contents: Vec<u8>,
}

/// Session encapsulates a UDP stream session
struct Session {
    log: Logger,
    send: SendHalf,
    /// dest is where to send data to
    dest: SocketAddr,
    /// from is the original sender
    from: SocketAddr,
    // TODO: store a session expiry, and update when you send data
}

impl Session {
    // TODO: write some tests
    async fn new(
        base: &Logger,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<Self> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let (recv, send) = UdpSocket::bind(addr).await?.split();
        let mut s = Session {
            log: base.new(o!("source" => "server::Session", "from" => from, "dest" => dest)),
            send,
            from,
            dest,
        };
        debug!(s.log, "Session created");

        s.run(recv, sender);

        return Ok(s);
    }

    // TODO: write some tests
    /// run starts processing receiving udp packets on its UdpSocket
    fn run(&mut self, mut recv: RecvHalf, mut sender: Sender<Packet>) {
        let log = self.log.clone();
        let from = self.from;
        tokio::spawn(async move {
            let mut buf: Vec<u8> = vec![0; 65535];
            // TODO: work out how to shut this down once this session expires
            loop {
                debug!(log, "awaiting incoming packet");
                match recv.recv_from(&mut buf).await {
                    Err(err) => error!(log, "Error receiving packet: {}", err),
                    Ok((size, recv_addr)) => {
                        let packet = &buf[..size];
                        debug!(
                            log,
                            "Received packet from {}, {}",
                            recv_addr,
                            from_utf8(packet).unwrap()
                        );
                        if let Err(err) = sender
                            .send(Packet {
                                contents: packet.to_vec(),
                                dest: from,
                            })
                            .await
                        {
                            println!("Error sending packet to channel: {}", err);
                        }
                    }
                }
            }
        });
    }

    /// Both sends a packet to the Session's dest.
    async fn send_to(&mut self, buf: &[u8]) -> Result<usize> {
        debug!(
            self.log,
            "Sending packet to: {}, {}",
            self.dest,
            from_utf8(buf).unwrap()
        );
        return self.send.send_to(buf, &self.dest).await;
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
    use std::str::from_utf8;
    use std::sync::Arc;

    use tokio::net::UdpSocket;
    use tokio::sync::mpsc::channel;
    use tokio::sync::{oneshot, RwLock};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::logger;
    use crate::server::{Packet, Server, Session, SessionMap};

    #[tokio::test]
    async fn server_bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 12345);
        assert_eq!(expected, addr)
    }

    #[tokio::test]
    async fn server_process_receive_socket() {
        let log = logger();
        let (local_addr, wait) = assert_recv_udp().await;

        let config = Arc::new(Config {
            local: Local { port: 0 },
            connections: ConnectionConfig::Sender {
                address: local_addr,
                connection_id: String::from(""),
            },
        });
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let receive_socket = UdpSocket::bind(addr).await.unwrap();
        let receive_addr = receive_socket.local_addr().unwrap();
        let (mut recv, mut send) = receive_socket.split();
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, mut recv_packets) = channel::<Packet>(1);

        tokio::spawn(async move {
            Server::process_receive_socket(
                log,
                config,
                &mut recv,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
        });

        send.send_to("hello".as_bytes(), &receive_addr)
            .await
            .unwrap();

        wait.await.unwrap();
        recv_packets.close();
    }

    #[tokio::test]
    async fn server_ensure_session() {
        let log = logger();
        let map: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let from: SocketAddr = "127.0.0.1:27890".parse().unwrap();
        let dest: SocketAddr = "127.0.0.1:27891".parse().unwrap();
        let (sender, mut recv) = channel::<Packet>(1);

        // gate
        {
            assert!(map.read().await.is_empty());
        }
        Server::ensure_session(&log, map.clone(), from, dest, sender)
            .await
            .unwrap();

        let rmap = map.read().await;
        let key = (from, dest);
        assert!(rmap.contains_key(&key));

        let sess = rmap.get(&key).unwrap().lock().await;
        assert_eq!(from, sess.from);
        assert_eq!(dest, sess.dest);
        assert_eq!(1, rmap.keys().len());

        recv.close();
    }

    #[tokio::test]
    async fn server_process_receive_packet_channel() {
        let log = logger();
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();

        let (mut recv_socket, send_socket) = socket.split();
        let (mut send_packet, recv_packet) = channel::<Packet>(5);
        let (done, wait) = oneshot::channel::<()>();

        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv_socket.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });

        if let Err(err) = send_packet
            .send(Packet {
                dest: local_addr,
                contents: String::from("hello").into_bytes(),
            })
            .await
        {
            assert!(false, err)
        }

        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, recv_packet).await;
        });
        wait.await.unwrap();
    }

    #[tokio::test]
    async fn session_send_to() {
        let log = logger();
        let (sender, _) = channel::<Packet>(1);
        let (local_addr, wait) = assert_recv_udp().await;

        let mut session = Session::new(&log, local_addr, local_addr, sender)
            .await
            .unwrap();
        session.send_to("hello".as_bytes()).await.unwrap();
        wait.await.unwrap();
    }

    /// assert_recv_udp asserts that the returned SockerAddr recieved a UDP packet
    /// with the contents of "hello"
    /// call wait.await.unwrap() to see if the message was received
    async fn assert_recv_udp() -> (SocketAddr, oneshot::Receiver<()>) {
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();
        let (mut recv, _) = socket.split();
        let (done, wait) = oneshot::channel::<()>();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });
        (local_addr, wait)
    }

    /// ephemeral_socket provides a socket bound to an ephemeral port
    async fn ephemeral_socket() -> UdpSocket {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        UdpSocket::bind(addr).await.unwrap()
    }
}
<commit_msg>Tests for Session::new (#21)<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, error, info, o, Logger};
use tokio::io::Result;
use tokio::net::udp::{RecvHalf, SendHalf};
use tokio::net::UdpSocket;
use tokio::sync::mpsc::{channel, Receiver, Sender};
use tokio::sync::{Mutex, RwLock};

use crate::config::{Config, ConnectionConfig};

type SessionMap = Arc<RwLock<HashMap<(SocketAddr, SocketAddr), Mutex<Session>>>>;

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    // TODO: write tests for this
    /// start the async processing of incoming UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        self.log_config(&config);

        let (mut receive_socket, send_socket) = Server::bind(&config).await?.split();
        // HashMap key is from,destination addresses as a tuple.
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, receive_packets) = channel::<Packet>(1024);

        let log = self.log.clone();
        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, receive_packets).await;
            debug!(log, "Receiver closed");
        });

        loop {
            if let Err(err) = Server::process_receive_socket(
                self.log.clone(),
                config.clone(),
                &mut receive_socket,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
            {
                error!(self.log, "Error processing reeive socket: {}", err);
            }
        }
    }

    /// process_receive_socket takes packets from the local socket and asynchronously
    /// processes them to send them out to endpoints.
    async fn process_receive_socket(
        log: Logger,
        config: Arc<Config>,
        receive_socket: &mut RecvHalf,
        sessions: SessionMap,
        send_packets: Sender<Packet>,
    ) -> Result<()> {
        let mut buf: Vec<u8> = vec![0; 65535];
        let (size, recv_addr) = receive_socket.recv_from(&mut buf).await?;
        tokio::spawn(async move {
            let endpoints = config.get_endpoints();
            let packet = &buf[..size];

            debug!(
                log,
                "Packet Received from: {}, {}",
                recv_addr,
                from_utf8(packet).unwrap()
            );

            for (_, dest) in endpoints.iter() {
                if let Err(err) = Server::ensure_session(
                    &log,
                    sessions.clone(),
                    recv_addr,
                    *dest,
                    send_packets.clone(),
                )
                .await
                {
                    error!(log, "Error ensuring session exists: {}", err);
                }

                let map = sessions.read().await;
                let mut session = map.get(&(recv_addr, *dest)).unwrap().lock().await;
                if let Err(err) = session.send_to(packet).await {
                    error!(log, "Error ensuring sending packet: {}", err)
                }
            }
        });
        return Ok(());
    }

    /// process_receive_packet_channel blocks on receive_packets.recv() channel
    /// and sends each packet on to the Packet.dest
    async fn process_receive_packet_channel(
        log: &Logger,
        mut send_socket: SendHalf,
        mut receive_packets: Receiver<Packet>,
    ) {
        while let Some(packet) = receive_packets.recv().await {
            debug!(
                log,
                "Sending packet back to origin";
                "origin" => packet.dest,
                "contents" => String::from_utf8(packet.contents.clone()).unwrap(),
            );

            if let Err(err) = send_socket
                .send_to(packet.contents.as_slice(), &packet.dest)
                .await
            {
                error!(log, "Error sending packet"; "dest" => packet.dest.to_string(), "err" => err.to_string());
            }
        }
    }

    /// log_config outputs a log of what is configured
    fn log_config(&self, config: &Arc<Config>) {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// ensure_session makes sure there is a value session for the name in the sessions map
    async fn ensure_session(
        log: &Logger,
        sessions: SessionMap,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<()> {
        {
            let map = sessions.read().await;
            if map.contains_key(&(from, dest)) {
                return Ok(());
            }
        }
        let s = Session::new(log, from, dest, sender).await?;
        {
            let mut map = sessions.write().await;
            map.insert((from, dest), Mutex::new(s));
        }
        return Ok(());
    }
}

/// Packet represents a packet that needs to go somewhere
struct Packet {
    dest: SocketAddr,
    contents: Vec<u8>,
}

/// Session encapsulates a UDP stream session
struct Session {
    log: Logger,
    send: SendHalf,
    /// dest is where to send data to
    dest: SocketAddr,
    /// from is the original sender
    from: SocketAddr,
    // TODO: store a session expiry, and update when you send data
}

impl Session {
    /// new creates a new Session, and starts the process of receiving udp sockets
    /// from its ephemeral port from endpoint(s)
    async fn new(
        base: &Logger,
        from: SocketAddr,
        dest: SocketAddr,
        sender: Sender<Packet>,
    ) -> Result<Self> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let (recv, send) = UdpSocket::bind(addr).await?.split();
        let mut s = Session {
            log: base.new(o!("source" => "server::Session", "from" => from, "dest" => dest)),
            send,
            from,
            dest,
        };
        debug!(s.log, "Session created");

        s.run(recv, sender);
        Ok(s)
    }

    /// run starts processing receiving udp packets on its UdpSocket
    fn run(&mut self, mut recv: RecvHalf, mut sender: Sender<Packet>) {
        let log = self.log.clone();
        let from = self.from;
        tokio::spawn(async move {
            let mut buf: Vec<u8> = vec![0; 65535];
            // TODO: work out how to shut this down once this session expires
            loop {
                debug!(log, "awaiting incoming packet");
                match recv.recv_from(&mut buf).await {
                    Err(err) => error!(log, "Error receiving packet: {}", err),
                    Ok((size, recv_addr)) => {
                        let packet = &buf[..size];
                        debug!(
                            log,
                            "Received packet from {}, {}",
                            recv_addr,
                            from_utf8(packet).unwrap()
                        );
                        if let Err(err) = sender
                            .send(Packet {
                                contents: packet.to_vec(),
                                dest: from,
                            })
                            .await
                        {
                            println!("Error sending packet to channel: {}", err);
                        }
                    }
                }
            }
        });
    }

    /// Sends a packet to the Session's dest.
    async fn send_to(&mut self, buf: &[u8]) -> Result<usize> {
        debug!(
            self.log,
            "Sending packet to: {}, {}",
            self.dest,
            from_utf8(buf).unwrap()
        );
        return self.send.send_to(buf, &self.dest).await;
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
    use std::str::from_utf8;
    use std::sync::Arc;

    use tokio::net::UdpSocket;
    use tokio::sync::mpsc::channel;
    use tokio::sync::{oneshot, RwLock};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::logger;
    use crate::server::{Packet, Server, Session, SessionMap};

    #[tokio::test]
    async fn server_bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 12345);
        assert_eq!(expected, addr)
    }

    #[tokio::test]
    async fn server_process_receive_socket() {
        let log = logger();
        let (local_addr, wait) = assert_recv_udp().await;

        let config = Arc::new(Config {
            local: Local { port: 0 },
            connections: ConnectionConfig::Sender {
                address: local_addr,
                connection_id: String::from(""),
            },
        });
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        let receive_socket = UdpSocket::bind(addr).await.unwrap();
        let receive_addr = receive_socket.local_addr().unwrap();
        let (mut recv, mut send) = receive_socket.split();
        let sessions: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let (send_packets, mut recv_packets) = channel::<Packet>(1);

        tokio::spawn(async move {
            Server::process_receive_socket(
                log,
                config,
                &mut recv,
                sessions.clone(),
                send_packets.clone(),
            )
            .await
        });

        send.send_to("hello".as_bytes(), &receive_addr)
            .await
            .unwrap();

        wait.await.unwrap();
        recv_packets.close();
    }

    #[tokio::test]
    async fn server_ensure_session() {
        let log = logger();
        let map: SessionMap = Arc::new(RwLock::new(HashMap::new()));
        let from: SocketAddr = "127.0.0.1:27890".parse().unwrap();
        let dest: SocketAddr = "127.0.0.1:27891".parse().unwrap();
        let (sender, mut recv) = channel::<Packet>(1);

        // gate
        {
            assert!(map.read().await.is_empty());
        }
        Server::ensure_session(&log, map.clone(), from, dest, sender)
            .await
            .unwrap();

        let rmap = map.read().await;
        let key = (from, dest);
        assert!(rmap.contains_key(&key));

        let sess = rmap.get(&key).unwrap().lock().await;
        assert_eq!(from, sess.from);
        assert_eq!(dest, sess.dest);
        assert_eq!(1, rmap.keys().len());

        recv.close();
    }

    #[tokio::test]
    async fn server_process_receive_packet_channel() {
        let log = logger();
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();

        let (mut recv_socket, send_socket) = socket.split();
        let (mut send_packet, recv_packet) = channel::<Packet>(5);
        let (done, wait) = oneshot::channel::<()>();

        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv_socket.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });

        if let Err(err) = send_packet
            .send(Packet {
                dest: local_addr,
                contents: String::from("hello").into_bytes(),
            })
            .await
        {
            assert!(false, err)
        }

        tokio::spawn(async move {
            Server::process_receive_packet_channel(&log, send_socket, recv_packet).await;
        });
        wait.await.unwrap();
    }

    #[tokio::test]
    async fn session_new() {
        let log = logger();
        let mut socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();
        let (send_packet, mut recv_packet) = channel::<Packet>(5);

        let mut sess = Session::new(&log, local_addr, local_addr, send_packet)
            .await
            .unwrap();

        // echo the packet back again
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let (size, recv_addr) = socket.recv_from(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            socket.send_to(&buf[..size], recv_addr).await.unwrap();
        });

        sess.send_to("hello".as_bytes()).await.unwrap();

        let packet = recv_packet.recv().await.unwrap();
        assert_eq!(String::from("hello").into_bytes(), packet.contents);
        assert_eq!(local_addr, packet.dest);
    }

    #[tokio::test]
    async fn session_send_to() {
        let log = logger();
        let (sender, _) = channel::<Packet>(1);
        let (local_addr, wait) = assert_recv_udp().await;

        let mut session = Session::new(&log, local_addr, local_addr, sender)
            .await
            .unwrap();
        session.send_to("hello".as_bytes()).await.unwrap();
        wait.await.unwrap();
    }

    /// assert_recv_udp asserts that the returned SockerAddr recieved a UDP packet
    /// with the contents of "hello"
    /// call wait.await.unwrap() to see if the message was received
    async fn assert_recv_udp() -> (SocketAddr, oneshot::Receiver<()>) {
        let socket = ephemeral_socket().await;
        let local_addr = socket.local_addr().unwrap();
        let (mut recv, _) = socket.split();
        let (done, wait) = oneshot::channel::<()>();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = recv.recv(&mut buf).await.unwrap();
            assert_eq!("hello", from_utf8(&buf[..size]).unwrap());
            done.send(()).unwrap();
        });
        (local_addr, wait)
    }

    /// ephemeral_socket provides a socket bound to an ephemeral port
    async fn ephemeral_socket() -> UdpSocket {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        UdpSocket::bind(addr).await.unwrap()
    }
}
<|endoftext|>


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

#[macro_export]
macro_rules! connect {
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $ty:ident :: $method:ident ( $($arg:expr),* )) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method($($arg),*)
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $ty:ident :: $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )*, $target:expr, $ty:ident :: $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move || {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $method:ident ( $($arg:expr),* )) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut Self = unsafe { &mut *target };
            target.$method($($arg),*)
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut Self = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )*, $target:expr, $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move || {
            let target: &mut Self = unsafe { &mut *target };
            target.$method()
        });
    };
}
<commit_msg>Added connect_static macro.<commit_after>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#[macro_export]
macro_rules! connect {
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $ty:ident :: $method:ident ( $($arg:expr),* )) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method($($arg),*)
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $ty:ident :: $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )*, $target:expr, $ty:ident :: $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move || {
            let target: &mut $ty = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $method:ident ( $($arg:expr),* )) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut Self = unsafe { &mut *target };
            target.$method($($arg),*)
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move |$($ident),*| {
            let target: &mut Self = unsafe { &mut *target };
            target.$method()
        });
    };
    ($source:expr, $event:ident $( [ $($param:expr),* ] )*, $target:expr, $method:ident) => {
        let target = &mut *$target as *mut _;
        $source.$event($($($param,)*)* move || {
            let target: &mut Self = unsafe { &mut *target };
            target.$method()
        });
    };
}

#[macro_export]
macro_rules! connect_static {
    ($source:ident, $event:ident $( [ $($param:expr),* ] )* ( $($ident:pat),* ), $target:expr, $method:ident ( $($arg:expr),* )) => {
        let target = &mut *$target as *mut _;
        $source::$event($($($param,)*)* move |$($ident),*| {
            let target: &mut Self = unsafe { &mut *target };
            target.$method($($arg),*)
        });
    };
    ($source:ident, $event:ident $( [ $($param:expr),* ] )*, $target:expr, $method:ident) => {
        let target = &mut *$target as *mut _;
        $source::$event($($($param,)*)* move || {
            let target: &mut Self = unsafe { &mut *target };
            target.$method()
        });
    };
}
<|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};
#[cfg(feature = "evented")]
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 ffi::{self,NlFamily,GenlId,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,NlmF};
use genlhdr::GenlHdr;
use nlattr::NlAttrHdr;
use nlhdr::NlHdr;

/// Handle for the socket file descriptor
#[cfg(feature = "evented")]
pub struct NlSocket<T, P> {
    fd: c_int,
    poll: mio::Poll,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
}

/// Handle for the socket file descriptor
#[cfg(not(feature = "evented"))]
pub struct NlSocket<T, P> {
    fd: c_int,
    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
    #[cfg(feature = "evented")]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, 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)
    }

    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(not(feature = "evented"))]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, 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::AF_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::nlhdr::NlHdr`)
    pub fn send<B>(&mut self, buf: StreamReadBuffer<B>, flags: i32) -> Result<isize, 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 `NlHdr` struct
    pub fn send_nl(&mut self, msg: NlHdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(StreamReadBuffer::new(mem), 0)?;
        Ok(())
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a>(&mut self, mut buf: StreamWriteBuffer<'a>, flags: i32)
            -> Result<StreamReadBuffer<StreamWriteBuffer<'a>>, io::Error> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => {
                buf.set_bytes_written(i as usize)?;
                Ok(StreamReadBuffer::new(buf))
            },
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Convenience function to receive an `NlHdr` struct
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<NlHdr<T, P>, NlError> {
        let mem_write = StreamWriteBuffer::new_growable(buf_sz.or(Some(MAX_NL_LENGTH)));
        let mut mem_read = self.recv(mem_write, 0)?;
        Ok(NlHdr::<T, P>::deserialize(&mut mem_read)?)
    }

    /// Convenience function to receive an `NlHdr` struct with function type parameters
    /// that determine deserialization type
    pub fn recv_nl_typed<TT, PP>(&mut self, buf_sz: Option<usize>)
            -> Result<NlHdr<TT, PP>, NlError> where TT: Nl + Into<u16> + From<u16>, PP: Nl {
        let mem_write = StreamWriteBuffer::new_growable(buf_sz.or(Some(MAX_NL_LENGTH)));
        let mut mem_read = self.recv(mem_write, 0)?;
        Ok(NlHdr::<TT, PP>::deserialize(&mut mem_read)?)
    }

    /// 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::<ffi::Nlmsg, Nlmsgerr<ffi::Nlmsg>>(buf_sz)?;
        if ack.nl_type == ffi::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, GenlHdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, GenlHdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<NlHdr<GenlId, GenlHdr<CtrlCmd>>, NlError> {
        let attrs = vec![NlAttrHdr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = GenlHdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = NlHdr::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 = NlHdr<T, P>;
    type Error = ();

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        if !self.is_blocking().map_err(|_| ())? {
            return Err(());
        }
        let mut mem = StreamWriteBuffer::new_growable(Some(MAX_NL_LENGTH));
        let bytes_written = match self.poll_read(mem.as_mut()) {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(0)) => return Ok(Async::Ready(None)),
            Ok(Async::Ready(i)) => i,
            Err(_) => return Err(()),
        };
        mem.set_bytes_written(bytes_written);
        let mut mem_read = StreamReadBuffer::new(mem);
        let hdr = match NlHdr::<T, P>::deserialize(&mut mem_read) {
            Ok(h) => h,
            Err(_) => return Err(()),
        };
        Ok(Async::Ready(Some(hdr)))
    }
}

#[cfg(feature = "evented")]
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 ffi::{CtrlCmd,Nlmsg};
    use genlhdr::GenlHdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, GenlHdr<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>Fix unused error discard for stream feature<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};
#[cfg(feature = "evented")]
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 ffi::{self,NlFamily,GenlId,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,NlmF};
use genlhdr::GenlHdr;
use nlattr::NlAttrHdr;
use nlhdr::NlHdr;

/// Handle for the socket file descriptor
#[cfg(feature = "evented")]
pub struct NlSocket<T, P> {
    fd: c_int,
    poll: mio::Poll,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
}

/// Handle for the socket file descriptor
#[cfg(not(feature = "evented"))]
pub struct NlSocket<T, P> {
    fd: c_int,
    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
    #[cfg(feature = "evented")]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, 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)
    }

    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(not(feature = "evented"))]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, 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::AF_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::nlhdr::NlHdr`)
    pub fn send<B>(&mut self, buf: StreamReadBuffer<B>, flags: i32) -> Result<isize, 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 `NlHdr` struct
    pub fn send_nl(&mut self, msg: NlHdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(StreamReadBuffer::new(mem), 0)?;
        Ok(())
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a>(&mut self, mut buf: StreamWriteBuffer<'a>, flags: i32)
            -> Result<StreamReadBuffer<StreamWriteBuffer<'a>>, io::Error> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => {
                buf.set_bytes_written(i as usize)?;
                Ok(StreamReadBuffer::new(buf))
            },
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Convenience function to receive an `NlHdr` struct
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<NlHdr<T, P>, NlError> {
        let mem_write = StreamWriteBuffer::new_growable(buf_sz.or(Some(MAX_NL_LENGTH)));
        let mut mem_read = self.recv(mem_write, 0)?;
        Ok(NlHdr::<T, P>::deserialize(&mut mem_read)?)
    }

    /// Convenience function to receive an `NlHdr` struct with function type parameters
    /// that determine deserialization type
    pub fn recv_nl_typed<TT, PP>(&mut self, buf_sz: Option<usize>)
            -> Result<NlHdr<TT, PP>, NlError> where TT: Nl + Into<u16> + From<u16>, PP: Nl {
        let mem_write = StreamWriteBuffer::new_growable(buf_sz.or(Some(MAX_NL_LENGTH)));
        let mut mem_read = self.recv(mem_write, 0)?;
        Ok(NlHdr::<TT, PP>::deserialize(&mut mem_read)?)
    }

    /// 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::<ffi::Nlmsg, Nlmsgerr<ffi::Nlmsg>>(buf_sz)?;
        if ack.nl_type == ffi::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, GenlHdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, GenlHdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<NlHdr<GenlId, GenlHdr<CtrlCmd>>, NlError> {
        let attrs = vec![NlAttrHdr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = GenlHdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = NlHdr::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 = NlHdr<T, P>;
    type Error = ();

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        if !self.is_blocking().map_err(|_| ())? {
            return Err(());
        }
        let mut mem = StreamWriteBuffer::new_growable(Some(MAX_NL_LENGTH));
        let bytes_written = match self.poll_read(mem.as_mut()) {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(0)) => return Ok(Async::Ready(None)),
            Ok(Async::Ready(i)) => i,
            Err(_) => return Err(()),
        };
        match mem.set_bytes_written(bytes_written) {
            Ok(()) => (),
            Err(_) => return Err(()),
        };
        let mut mem_read = StreamReadBuffer::new(mem);
        let hdr = match NlHdr::<T, P>::deserialize(&mut mem_read) {
            Ok(h) => h,
            Err(_) => return Err(()),
        };
        Ok(Async::Ready(Some(hdr)))
    }
}

#[cfg(feature = "evented")]
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 ffi::{CtrlCmd,Nlmsg};
    use genlhdr::GenlHdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, GenlHdr<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 super::*;
use std::fmt::Debug;

#[test]
fn stream_basic()
{
    let sink = Sink::new();
    let stream = sink.stream();
    let rx = stream.channel();
    let signal = stream.hold(0);

    sink.send(42);
    sink.send(33);
    sink.feed(0..5);

    let result: Vec<i32> = rx.try_iter().collect();
    assert_eq!(result, [42, 33, 0, 1, 2, 3, 4]);
    assert_eq!(signal.sample(), 4);
}

#[test]
fn signal_basic()
{
    let signal = Signal::constant(42);
    assert_eq!(signal.sample(), 42);
    signal.sample_with(|val| assert_eq!(*val, 42));

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

    let sink = Sink::new();
    let stream = sink.stream();
    let s_double = stream.hold(0).map(|a| *a * 2);
    let s_odd = stream.hold_if(0, |a| *a % 2 != 0);

    sink.send(2);
    sink.send(3);
    sink.send(4);

    assert_eq!(s_double.sample(), 8);
    assert_eq!(s_odd.sample(), 3);
}

fn vec_cons<T: Clone>(mut v: Vec<T>, x: Cow<T>) -> Vec<T> { v.push(x.into_owned()); v }

#[test]
fn stream_operations()
{
    let sink: Sink<i32> = Sink::new();
    let stream = sink.stream();

    let s_string = stream.map(|a| a.to_string()).fold(vec![], vec_cons);
    let s_odd = stream.filter(|a| *a % 2 != 0).fold(vec![], vec_cons);
    let s_even_half = stream.filter_map(|a| if *a % 2 == 0 { Some(*a / 2) } else { None }).fold(vec![], vec_cons);
    let (pos, neg) = stream.map(|a| if *a > 0 { Ok(*a) } else { Err(*a) }).split();
    let s_pos = pos.fold(vec![], vec_cons);
    let s_neg = neg.fold(vec![], vec_cons);
    let s_merged = pos.merge(&neg.map(|a| -*a)).fold(vec![], vec_cons);
    let s_accum = stream.fold(vec![], vec_cons).snapshot(&stream, |s, _| s.into_owned()).fold(vec![], vec_cons);

    sink.feed(vec![5, 8, 13, -2, 42, -33]);

    assert_eq!(s_string.sample(), ["5", "8", "13", "-2", "42", "-33"]);
    assert_eq!(s_odd.sample(), [5, 13, -33]);
    assert_eq!(s_even_half.sample(), [4, -1, 21]);
    assert_eq!(s_pos.sample(), [5, 8, 13, 42]);
    assert_eq!(s_neg.sample(), [-2, -33]);
    assert_eq!(s_merged.sample(), [5, 8, 13, 2, 42, 33]);
    assert_eq!(s_accum.sample(), [vec![5], vec![5, 8], vec![5, 8, 13], vec![5, 8, 13, -2], vec![5, 8, 13, -2, 42], vec![5, 8, 13, -2, 42, -33]]);
}

#[cfg(feature="either")]
#[test]
fn merge_with()
{
    let sink1: Sink<i32> = Sink::new();
    let sink2: Sink<f32> = Sink::new();
    let stream: Stream<Result<_, _>> = sink1.stream().merge_with(&sink2.stream(), |e| e.either(|l| Ok(*l), |r| Err(*r)));
    let result = stream.fold(vec![], vec_cons);

    sink1.send(1);
    sink2.send(2.0);
    sink1.send(3);
    sink1.send(4);
    sink2.send(5.0);

    assert_eq!(result.sample(), [Ok(1), Err(2.0), Ok(3), Ok(4), Err(5.0)]);
}

#[test]
fn stream_channel()
{
    let sink = Sink::new();
    let rx = sink.stream().channel();

    let thread = std::thread::spawn(move || {
        let mut sum = 0;
        while let Some(val) = rx.recv().unwrap()
        {
            sum += val;
        }
        sum
    });
    sink.feed((1..100).map(Some));
    sink.send(None);

    let result = thread.join().unwrap();
    assert_eq!(result, 4950);
}

#[test]
fn stream_switch()
{
    let stream_sink = Sink::new();
    let sink1 = Sink::new();
    let sink2 = Sink::new();

    let switched = stream_sink.stream().switch();
    let events = switched.channel();

    sink1.send(1);
    sink2.send(2);

    stream_sink.send(sink2.stream());
    sink1.send(3);
    sink2.send(4);
    assert_eq!(events.try_recv(), Ok(4));

    stream_sink.send(sink1.stream());
    sink1.send(5);
    sink2.send(6);
    assert_eq!(events.try_recv(), Ok(5));
}

#[test]
fn signal_switch()
{
    let signal_sink = Sink::new();
    let switched = signal_sink.stream().hold(Signal::constant(0)).switch();

    signal_sink.send(Signal::constant(1));
    assert_eq!(switched.sample(), 1);

    signal_sink.send(Signal::constant(2));
    assert_eq!(switched.sample(), 2);
}

#[derive(Debug)]
struct Storage<T>(Vec<T>);

impl<T> Storage<T>
{
    fn new() -> Self { Storage(Vec::new()) }
    fn push(mut self, a: T) -> Self { self.0.push(a); self }
}

impl<T: Debug> Clone for Storage<T>
{
    fn clone(&self) -> Self { panic!("storage cloned! {:?}", self.0) }
}

#[test]
fn cloning()
{
    let sink = Sink::new();
    let accum = sink.stream().fold(Storage::new(), |a, v| a.push(*v));

    sink.feed(0..5);
    accum.sample_with(|res| assert_eq!(res.0, [0, 1, 2, 3, 4]));
}

#[test]
fn filter_extra()
{
    let sink = Sink::new();
    let stream = sink.stream();
    let sign_res = stream.map(|a| if *a >= 0 { Ok(*a) } else { Err(*a) });
    let even_opt = stream.map(|a| if *a % 2 == 0 { Some(*a) } else { None });
    let s_even = even_opt.filter_some().fold(vec![], vec_cons);
    let s_pos = sign_res.filter_first().fold(vec![], vec_cons);
    let s_neg = sign_res.filter_second().fold(vec![], vec_cons);

    sink.feed(vec![1, 8, -3, 42, -66]);

    assert_eq!(s_even.sample(), [8, 42, -66]);
    assert_eq!(s_pos.sample(), [1, 8, 42]);
    assert_eq!(s_neg.sample(), [-3, -66]);
}
<commit_msg>make the stream_channel test a bit more complex<commit_after>use super::*;
use std::fmt::Debug;

#[test]
fn stream_basic()
{
    let sink = Sink::new();
    let stream = sink.stream();
    let rx = stream.channel();
    let signal = stream.hold(0);

    sink.send(42);
    sink.send(33);
    sink.feed(0..5);

    let result: Vec<i32> = rx.try_iter().collect();
    assert_eq!(result, [42, 33, 0, 1, 2, 3, 4]);
    assert_eq!(signal.sample(), 4);
}

#[test]
fn signal_basic()
{
    let signal = Signal::constant(42);
    assert_eq!(signal.sample(), 42);
    signal.sample_with(|val| assert_eq!(*val, 42));

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

    let sink = Sink::new();
    let stream = sink.stream();
    let s_double = stream.hold(0).map(|a| *a * 2);
    let s_odd = stream.hold_if(0, |a| *a % 2 != 0);

    sink.send(2);
    sink.send(3);
    sink.send(4);

    assert_eq!(s_double.sample(), 8);
    assert_eq!(s_odd.sample(), 3);
}

fn vec_cons<T: Clone>(mut v: Vec<T>, x: Cow<T>) -> Vec<T> { v.push(x.into_owned()); v }

#[test]
fn stream_operations()
{
    let sink: Sink<i32> = Sink::new();
    let stream = sink.stream();

    let s_string = stream.map(|a| a.to_string()).fold(vec![], vec_cons);
    let s_odd = stream.filter(|a| *a % 2 != 0).fold(vec![], vec_cons);
    let s_even_half = stream.filter_map(|a| if *a % 2 == 0 { Some(*a / 2) } else { None }).fold(vec![], vec_cons);
    let (pos, neg) = stream.map(|a| if *a > 0 { Ok(*a) } else { Err(*a) }).split();
    let s_pos = pos.fold(vec![], vec_cons);
    let s_neg = neg.fold(vec![], vec_cons);
    let s_merged = pos.merge(&neg.map(|a| -*a)).fold(vec![], vec_cons);
    let s_accum = stream.fold(vec![], vec_cons).snapshot(&stream, |s, _| s.into_owned()).fold(vec![], vec_cons);

    sink.feed(vec![5, 8, 13, -2, 42, -33]);

    assert_eq!(s_string.sample(), ["5", "8", "13", "-2", "42", "-33"]);
    assert_eq!(s_odd.sample(), [5, 13, -33]);
    assert_eq!(s_even_half.sample(), [4, -1, 21]);
    assert_eq!(s_pos.sample(), [5, 8, 13, 42]);
    assert_eq!(s_neg.sample(), [-2, -33]);
    assert_eq!(s_merged.sample(), [5, 8, 13, 2, 42, 33]);
    assert_eq!(s_accum.sample(), [vec![5], vec![5, 8], vec![5, 8, 13], vec![5, 8, 13, -2], vec![5, 8, 13, -2, 42], vec![5, 8, 13, -2, 42, -33]]);
}

#[cfg(feature="either")]
#[test]
fn merge_with()
{
    let sink1: Sink<i32> = Sink::new();
    let sink2: Sink<f32> = Sink::new();
    let stream: Stream<Result<_, _>> = sink1.stream().merge_with(&sink2.stream(), |e| e.either(|l| Ok(*l), |r| Err(*r)));
    let result = stream.fold(vec![], vec_cons);

    sink1.send(1);
    sink2.send(2.0);
    sink1.send(3);
    sink1.send(4);
    sink2.send(5.0);

    assert_eq!(result.sample(), [Ok(1), Err(2.0), Ok(3), Ok(4), Err(5.0)]);
}

#[test]
fn stream_channel()
{
    let sink = Sink::new();
    let rx = sink.stream().channel();

    let thread = std::thread::spawn(move || {
        let sink2 = Sink::new();
        let s_sum = sink2.stream().fold(0, |a, n| a + *n);
        while let Ok(_) = rx.recv().map(|v| sink2.send(v)) { /* empty */ }
        s_sum.sample()
    });

    sink.feed(1..100);
    drop(sink);

    let result = thread.join().unwrap();
    assert_eq!(result, 4950);
}

#[test]
fn stream_switch()
{
    let stream_sink = Sink::new();
    let sink1 = Sink::new();
    let sink2 = Sink::new();

    let switched = stream_sink.stream().switch();
    let events = switched.channel();

    sink1.send(1);
    sink2.send(2);

    stream_sink.send(sink2.stream());
    sink1.send(3);
    sink2.send(4);
    assert_eq!(events.try_recv(), Ok(4));

    stream_sink.send(sink1.stream());
    sink1.send(5);
    sink2.send(6);
    assert_eq!(events.try_recv(), Ok(5));
}

#[test]
fn signal_switch()
{
    let signal_sink = Sink::new();
    let switched = signal_sink.stream().hold(Signal::constant(0)).switch();

    signal_sink.send(Signal::constant(1));
    assert_eq!(switched.sample(), 1);

    signal_sink.send(Signal::constant(2));
    assert_eq!(switched.sample(), 2);
}

#[derive(Debug)]
struct Storage<T>(Vec<T>);

impl<T> Storage<T>
{
    fn new() -> Self { Storage(Vec::new()) }
    fn push(mut self, a: T) -> Self { self.0.push(a); self }
}

impl<T: Debug> Clone for Storage<T>
{
    fn clone(&self) -> Self { panic!("storage cloned! {:?}", self.0) }
}

#[test]
fn cloning()
{
    let sink = Sink::new();
    let accum = sink.stream().fold(Storage::new(), |a, v| a.push(*v));

    sink.feed(0..5);
    accum.sample_with(|res| assert_eq!(res.0, [0, 1, 2, 3, 4]));
}

#[test]
fn filter_extra()
{
    let sink = Sink::new();
    let stream = sink.stream();
    let sign_res = stream.map(|a| if *a >= 0 { Ok(*a) } else { Err(*a) });
    let even_opt = stream.map(|a| if *a % 2 == 0 { Some(*a) } else { None });
    let s_even = even_opt.filter_some().fold(vec![], vec_cons);
    let s_pos = sign_res.filter_first().fold(vec![], vec_cons);
    let s_neg = sign_res.filter_second().fold(vec![], vec_cons);

    sink.feed(vec![1, 8, -3, 42, -66]);

    assert_eq!(s_even.sample(), [8, 42, -66]);
    assert_eq!(s_pos.sample(), [1, 8, 42]);
    assert_eq!(s_neg.sample(), [-3, -66]);
}
<|endoftext|>


<commit_before>use ::CanSocket;

#[test]
fn test_nonexistant_device() {
    assert!(CanSocket::open("invalid").is_err());
}

#[cfg(feature = "vcan_tests")]
mod vcan_tests {
    use ::{CanSocket, ERR_MASK_ALL, ERR_MASK_NONE};
    use std::time;
    use ::ShouldRetry;

    #[test]
    fn vcan0_timeout() {
        let cs = CanSocket::open("vcan0").unwrap();
        cs.set_read_timeout(time::Duration::from_millis(100)).unwrap();
        assert!(cs.read_frame().should_retry());
    }

    #[test]
    fn vcan0_set_error_mask() {
        let cs = CanSocket::open("vcan0").unwrap();
        cs.set_error_mask(ERR_MASK_ALL).unwrap();
        cs.set_error_mask(ERR_MASK_NONE).unwrap();
    }
}
<commit_msg>Add missing imports in tests.<commit_after>use ::CanSocket;

#[test]
fn test_nonexistant_device() {
    assert!(CanSocket::open("invalid").is_err());
}

#[cfg(feature = "vcan_tests")]
mod vcan_tests {
    use ::{CanFrame, CanSocket, ERR_MASK_ALL, ERR_MASK_NONE};
    use std::time;
    use ::ShouldRetry;

    #[test]
    fn vcan0_timeout() {
        let cs = CanSocket::open("vcan0").unwrap();
        cs.set_read_timeout(time::Duration::from_millis(100)).unwrap();
        assert!(cs.read_frame().should_retry());
    }

    #[test]
    fn vcan0_set_error_mask() {
        let cs = CanSocket::open("vcan0").unwrap();
        cs.set_error_mask(ERR_MASK_ALL).unwrap();
        cs.set_error_mask(ERR_MASK_NONE).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::collections;
use std::fmt;
use std::ops;
use std::slice;

/// A sentence consists of zero or more `Token`s. It is a small wrapper
/// around `Vec<Token>` that provides some extra convenience. For example,
/// the implementation of the `Display` trait outputs the sentence in
/// CoNLL-X format.
#[derive(Debug,PartialEq)]
pub struct Sentence {
    tokens: Vec<Token>,
}

impl Sentence {
    /// Create a new `Sentence` from a vector of tokens.
    pub fn new(tokens: Vec<Token>) -> Sentence {
        Sentence { tokens: tokens }
    }

    /// Get the underlying vector of tokens.
    pub fn as_tokens(&self) -> &Vec<Token> {
        &self.tokens
    }

    /// Get an iterator over the sentence tokens.
    pub fn iter(&self) -> slice::Iter<Token> {
        self.tokens.iter()
    }

    /// Get a mutable iterator over the sentence tokens.
    pub fn iter_mut(&mut self) -> slice::IterMut<Token> {
        self.tokens.iter_mut()
    }
}

impl fmt::Display for Sentence {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let last_id = self.tokens.len() - 1;

        for (id, token) in self.iter().enumerate() {
            if id == last_id {
                try!(write!(f, "{}\t{}", id + 1, token))
            } else {
                try!(write!(f, "{}\t{}\n", id + 1, token))
            }
        }

        Ok(())
    }
}

impl ops::Index<usize> for Sentence {
    type Output = Token;
    fn index(&self, index: usize) -> &Token {
        &self.tokens[index]
    }
}

impl ops::IndexMut<usize> for Sentence {
    fn index_mut(&mut self, index: usize) -> &mut Token {
        &mut self.tokens[index]
    }
}

impl<'a> IntoIterator for &'a Sentence {
    type Item = &'a Token;
    type IntoIter = slice::Iter<'a, Token>;

    fn into_iter(self) -> slice::Iter<'a, Token> {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a mut Sentence {
    type Item = &'a mut Token;
    type IntoIter = slice::IterMut<'a, Token>;

    fn into_iter(mut self) -> slice::IterMut<'a, Token> {
        self.iter_mut()
    }
}

/// A builder for `Token`s.
///
/// The `Token` type stores a CoNLL-X token. However, since this format
/// permits a large number of fields, construction of a token can get
/// tedious. This builder provides a fluent interface for creating `Token`s.
pub struct TokenBuilder {
    token: Token,
}

impl TokenBuilder {
    /// Create a `Token` builder with all fields set to absent.
    pub fn new() -> TokenBuilder {
        TokenBuilder { token: Token::new() }
    }

    /// Set the word form or punctuation symbol.
    pub fn form<S>(mut self, form: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_form(Some(form));
        self
    }

    /// Set the lemma or stem of the word form.
    pub fn lemma<S>(mut self, lemma: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_lemma(Some(lemma));
        self
    }

    /// Set the coarse-grained part-of-speech tag.
    pub fn cpos<S>(mut self, cpos: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_cpos(Some(cpos));
        self
    }

    /// Set the fine-grained part-of-speech tag.
    pub fn pos<S>(mut self, pos: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_pos(Some(pos));
        self
    }

    /// Set the syntactic and/or morphological features of the token.
    pub fn features<S>(mut self, features: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_features(Some(features));
        self
    }

    /// Set the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn head(mut self, head: usize) -> TokenBuilder {
        self.token.set_head(Some(head));
        self
    }

    /// Set the dependency relation to the head of this token.
    pub fn head_rel<S>(mut self, head_rel: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_head_rel(Some(head_rel));
        self
    }

    /// Set the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn p_head(mut self, p_head: usize) -> TokenBuilder {
        self.token.set_p_head(Some(p_head));
        self
    }

    /// Set the dependency relation to the projective head of this token.
    pub fn p_head_rel<S>(mut self, p_head_rel: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_p_head_rel(Some(p_head_rel));
        self
    }

    pub fn token(self) -> Token {
        self.token
    }
}

/// A token in the CoNLL-X dependency format.
///
/// The fields of CoNLLX tokens are described at:
///
/// http://ilk.uvt.nl/conll/
///
/// This type provides exactly the same fields, except for the ID field
/// (since it can be derived from the sentence position of the token).
/// If a particular field is absent (*_* in the CoNLL-X format), its
/// value is `None`.
#[derive(Debug,PartialEq)]
pub struct Token {
    form: Option<String>,
    lemma: Option<String>,
    cpos: Option<String>,
    pos: Option<String>,
    features: Option<Features>,
    head: Option<usize>,
    head_rel: Option<String>,
    p_head: Option<usize>,
    p_head_rel: Option<String>,
}

impl Token {
    /// Create a new token where all the fields are absent.
    pub fn new() -> Token {
        Token {
            form: None,
            lemma: None,
            cpos: None,
            pos: None,
            features: None,
            head: None,
            head_rel: None,
            p_head: None,
            p_head_rel: None,
        }
    }

    /// Get the word form or punctuation symbol.
    pub fn form(&self) -> &Option<String> {
        &self.form
    }

    /// Get the lemma or stem of the word form.
    pub fn lemma(&self) -> &Option<String> {
        &self.lemma
    }

    /// Get the coarse-grained part-of-speech tag.
    pub fn cpos(&self) -> &Option<String> {
        &self.cpos
    }

    /// Get the fine-grained part-of-speech tag.
    pub fn pos(&self) -> &Option<String> {
        &self.pos
    }

    /// Get the syntactic and/or morphological features of the token.
    pub fn features(&self) -> &Option<Features> {
        &self.features
    }

    /// Get the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn head(&self) -> &Option<usize> {
        &self.head
    }

    /// Get the dependency relation to the head of this token.
    pub fn head_rel(&self) -> &Option<String> {
        &self.head_rel
    }

    /// Get the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn p_head(&self) -> &Option<usize> {
        &self.p_head
    }

    /// Get the dependency relation to the projective head of this token.
    pub fn p_head_rel(&self) -> &Option<String> {
        &self.p_head_rel
    }

    /// Set the word form or punctuation symbol.
    pub fn set_form<S>(&mut self, form: Option<S>)
        where S: Into<String>
    {
        self.form = form.map(|i| i.into())
    }

    /// Set the lemma or stem of the word form.
    pub fn set_lemma<S>(&mut self, lemma: Option<S>)
        where S: Into<String>
    {
        self.lemma = lemma.map(|i| i.into())
    }

    /// Set the coarse-grained part-of-speech tag.
    pub fn set_cpos<S>(&mut self, cpos: Option<S>)
        where S: Into<String>
    {
        self.cpos = cpos.map(|i| i.into())
    }

    /// Set the fine-grained part-of-speech tag.
    pub fn set_pos<S>(&mut self, pos: Option<S>)
        where S: Into<String>
    {
        self.pos = pos.map(|i| i.into())
    }

    /// Set the syntactic and/or morphological features of the token.
    pub fn set_features<S>(&mut self, features: Option<S>)
        where S: Into<String>
    {
        self.features = features.map(|s| Features { features: s.into() })
    }

    /// Set the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn set_head(&mut self, head: Option<usize>) {
        self.head = head
    }

    /// Set the dependency relation to the head of this token.
    pub fn set_head_rel<S>(&mut self, head_rel: Option<S>)
        where S: Into<String>
    {
        self.head_rel = head_rel.map(|i| i.into())
    }

    /// Set the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn set_p_head(&mut self, p_head: Option<usize>) {
        self.p_head = p_head
    }

    /// Set the dependency relation to the projective head of this token.
    pub fn set_p_head_rel<S>(&mut self, p_head_rel: Option<S>)
        where S: Into<String>
    {
        self.p_head_rel = p_head_rel.map(|i| i.into())
    }
}

/// This data type is used to represent features on a `Token`. In the CoNLL-X
/// specification, these are morphological features of the token. However,
/// they are used as a free form list or mapping of features.
#[derive(Debug,PartialEq)]
pub struct Features {
    features: String,
}

impl Features {
    /// Get the features field as a key-value mapping. This assumes that
    /// the key-value pairs are separed using a vertical bar (`|`) and keys
    /// and values using a colon (`:`). If the value is absent, corresponding
    /// value in the mapping is `None`.
    pub fn as_map(&self) -> collections::BTreeMap<String, Option<String>> {
        let mut features = collections::BTreeMap::new();

        for fv in self.features.split('|') {
            let mut iter = fv.split(':');
            if let Some(k) = iter.next() {
                let v = iter.next().map(|s| s.to_owned());
                features.insert(k.to_owned(), v.to_owned());
            }
        }

        features
    }

    /// Get the features field.
    pub fn as_str(&self) -> &str {
        self.features.as_ref()
    }
}

impl fmt::Display for Features {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.features.as_ref())
    }
}

pub const EMPTY_TOKEN: &'static str = "_";

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let head_str = self.head.as_ref().map(|n| n.to_string());
        let p_head_str = self.p_head.as_ref().map(|n| n.to_string());

        write!(f,
               "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}",
               self.form.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.lemma.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.cpos.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.pos.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.features.as_ref().map(|s| s.as_str()).unwrap_or(EMPTY_TOKEN),
               head_str.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.head_rel.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.p_head.clone().map(|n| n.to_string()).unwrap_or(EMPTY_TOKEN.to_string()),
               p_head_str.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN))
    }
}
<commit_msg>Modify Token getter return types.<commit_after>use std::collections;
use std::fmt;
use std::ops;
use std::slice;

/// A sentence consists of zero or more `Token`s. It is a small wrapper
/// around `Vec<Token>` that provides some extra convenience. For example,
/// the implementation of the `Display` trait outputs the sentence in
/// CoNLL-X format.
#[derive(Debug,PartialEq)]
pub struct Sentence {
    tokens: Vec<Token>,
}

impl Sentence {
    /// Create a new `Sentence` from a vector of tokens.
    pub fn new(tokens: Vec<Token>) -> Sentence {
        Sentence { tokens: tokens }
    }

    /// Get the underlying vector of tokens.
    pub fn as_tokens(&self) -> &Vec<Token> {
        &self.tokens
    }

    /// Get an iterator over the sentence tokens.
    pub fn iter(&self) -> slice::Iter<Token> {
        self.tokens.iter()
    }

    /// Get a mutable iterator over the sentence tokens.
    pub fn iter_mut(&mut self) -> slice::IterMut<Token> {
        self.tokens.iter_mut()
    }
}

impl fmt::Display for Sentence {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let last_id = self.tokens.len() - 1;

        for (id, token) in self.iter().enumerate() {
            if id == last_id {
                try!(write!(f, "{}\t{}", id + 1, token))
            } else {
                try!(write!(f, "{}\t{}\n", id + 1, token))
            }
        }

        Ok(())
    }
}

impl ops::Index<usize> for Sentence {
    type Output = Token;
    fn index(&self, index: usize) -> &Token {
        &self.tokens[index]
    }
}

impl ops::IndexMut<usize> for Sentence {
    fn index_mut(&mut self, index: usize) -> &mut Token {
        &mut self.tokens[index]
    }
}

impl<'a> IntoIterator for &'a Sentence {
    type Item = &'a Token;
    type IntoIter = slice::Iter<'a, Token>;

    fn into_iter(self) -> slice::Iter<'a, Token> {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a mut Sentence {
    type Item = &'a mut Token;
    type IntoIter = slice::IterMut<'a, Token>;

    fn into_iter(mut self) -> slice::IterMut<'a, Token> {
        self.iter_mut()
    }
}

/// A builder for `Token`s.
///
/// The `Token` type stores a CoNLL-X token. However, since this format
/// permits a large number of fields, construction of a token can get
/// tedious. This builder provides a fluent interface for creating `Token`s.
pub struct TokenBuilder {
    token: Token,
}

impl TokenBuilder {
    /// Create a `Token` builder with all fields set to absent.
    pub fn new() -> TokenBuilder {
        TokenBuilder { token: Token::new() }
    }

    /// Set the word form or punctuation symbol.
    pub fn form<S>(mut self, form: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_form(Some(form));
        self
    }

    /// Set the lemma or stem of the word form.
    pub fn lemma<S>(mut self, lemma: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_lemma(Some(lemma));
        self
    }

    /// Set the coarse-grained part-of-speech tag.
    pub fn cpos<S>(mut self, cpos: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_cpos(Some(cpos));
        self
    }

    /// Set the fine-grained part-of-speech tag.
    pub fn pos<S>(mut self, pos: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_pos(Some(pos));
        self
    }

    /// Set the syntactic and/or morphological features of the token.
    pub fn features<S>(mut self, features: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_features(Some(features));
        self
    }

    /// Set the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn head(mut self, head: usize) -> TokenBuilder {
        self.token.set_head(Some(head));
        self
    }

    /// Set the dependency relation to the head of this token.
    pub fn head_rel<S>(mut self, head_rel: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_head_rel(Some(head_rel));
        self
    }

    /// Set the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn p_head(mut self, p_head: usize) -> TokenBuilder {
        self.token.set_p_head(Some(p_head));
        self
    }

    /// Set the dependency relation to the projective head of this token.
    pub fn p_head_rel<S>(mut self, p_head_rel: S) -> TokenBuilder
        where S: Into<String>
    {
        self.token.set_p_head_rel(Some(p_head_rel));
        self
    }

    pub fn token(self) -> Token {
        self.token
    }
}

/// A token in the CoNLL-X dependency format.
///
/// The fields of CoNLLX tokens are described at:
///
/// http://ilk.uvt.nl/conll/
///
/// This type provides exactly the same fields, except for the ID field
/// (since it can be derived from the sentence position of the token).
/// If a particular field is absent (*_* in the CoNLL-X format), its
/// value is `None`.
#[derive(Debug,PartialEq)]
pub struct Token {
    form: Option<String>,
    lemma: Option<String>,
    cpos: Option<String>,
    pos: Option<String>,
    features: Option<Features>,
    head: Option<usize>,
    head_rel: Option<String>,
    p_head: Option<usize>,
    p_head_rel: Option<String>,
}

impl Token {
    /// Create a new token where all the fields are absent.
    pub fn new() -> Token {
        Token {
            form: None,
            lemma: None,
            cpos: None,
            pos: None,
            features: None,
            head: None,
            head_rel: None,
            p_head: None,
            p_head_rel: None,
        }
    }

    /// Get the word form or punctuation symbol.
    pub fn form(&self) -> Option<&str> {
        self.form.as_ref().map(String::as_ref)
    }

    /// Get the lemma or stem of the word form.
    pub fn lemma(&self) -> Option<&str> {
        self.lemma.as_ref().map(String::as_ref)
    }

    /// Get the coarse-grained part-of-speech tag.
    pub fn cpos(&self) -> Option<&str> {
        self.cpos.as_ref().map(String::as_ref)
    }

    /// Get the fine-grained part-of-speech tag.
    pub fn pos(&self) -> Option<&str> {
        self.pos.as_ref().map(String::as_ref)
    }

    /// Get the syntactic and/or morphological features of the token.
    pub fn features(&self) -> &Option<Features> {
        &self.features
    }

    /// Get the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn head(&self) -> Option<usize> {
        self.head
    }

    /// Get the dependency relation to the head of this token.
    pub fn head_rel(&self) -> Option<&str> {
        self.head_rel.as_ref().map(String::as_ref)
    }

    /// Get the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn p_head(&self) -> Option<usize> {
        self.p_head
    }

    /// Get the dependency relation to the projective head of this token.
    pub fn p_head_rel(&self) -> Option<&str> {
        self.p_head_rel.as_ref().map(String::as_ref)
    }

    /// Set the word form or punctuation symbol.
    pub fn set_form<S>(&mut self, form: Option<S>)
        where S: Into<String>
    {
        self.form = form.map(|i| i.into())
    }

    /// Set the lemma or stem of the word form.
    pub fn set_lemma<S>(&mut self, lemma: Option<S>)
        where S: Into<String>
    {
        self.lemma = lemma.map(|i| i.into())
    }

    /// Set the coarse-grained part-of-speech tag.
    pub fn set_cpos<S>(&mut self, cpos: Option<S>)
        where S: Into<String>
    {
        self.cpos = cpos.map(|i| i.into())
    }

    /// Set the fine-grained part-of-speech tag.
    pub fn set_pos<S>(&mut self, pos: Option<S>)
        where S: Into<String>
    {
        self.pos = pos.map(|i| i.into())
    }

    /// Set the syntactic and/or morphological features of the token.
    pub fn set_features<S>(&mut self, features: Option<S>)
        where S: Into<String>
    {
        self.features = features.map(|s| Features { features: s.into() })
    }

    /// Set the head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree.
    pub fn set_head(&mut self, head: Option<usize>) {
        self.head = head
    }

    /// Set the dependency relation to the head of this token.
    pub fn set_head_rel<S>(&mut self, head_rel: Option<S>)
        where S: Into<String>
    {
        self.head_rel = head_rel.map(|i| i.into())
    }

    /// Set the projective head of the token. This is the sentence position
    /// of the head **plus one**. If the head is 0, the token the root
    /// of the dependency tree. The dependency structure resulting from the
    /// projective heads must be projective.
    pub fn set_p_head(&mut self, p_head: Option<usize>) {
        self.p_head = p_head
    }

    /// Set the dependency relation to the projective head of this token.
    pub fn set_p_head_rel<S>(&mut self, p_head_rel: Option<S>)
        where S: Into<String>
    {
        self.p_head_rel = p_head_rel.map(|i| i.into())
    }
}

/// This data type is used to represent features on a `Token`. In the CoNLL-X
/// specification, these are morphological features of the token. However,
/// they are used as a free form list or mapping of features.
#[derive(Debug,PartialEq)]
pub struct Features {
    features: String,
}

impl Features {
    /// Get the features field as a key-value mapping. This assumes that
    /// the key-value pairs are separed using a vertical bar (`|`) and keys
    /// and values using a colon (`:`). If the value is absent, corresponding
    /// value in the mapping is `None`.
    pub fn as_map(&self) -> collections::BTreeMap<String, Option<String>> {
        let mut features = collections::BTreeMap::new();

        for fv in self.features.split('|') {
            let mut iter = fv.split(':');
            if let Some(k) = iter.next() {
                let v = iter.next().map(|s| s.to_owned());
                features.insert(k.to_owned(), v.to_owned());
            }
        }

        features
    }

    /// Get the features field.
    pub fn as_str(&self) -> &str {
        self.features.as_ref()
    }
}

impl fmt::Display for Features {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.features.as_ref())
    }
}

pub const EMPTY_TOKEN: &'static str = "_";

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let head_str = self.head.as_ref().map(|n| n.to_string());
        let p_head_str = self.p_head.as_ref().map(|n| n.to_string());

        write!(f,
               "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}",
               self.form.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.lemma.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.cpos.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.pos.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.features.as_ref().map(|s| s.as_str()).unwrap_or(EMPTY_TOKEN),
               head_str.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.head_rel.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN),
               self.p_head.clone().map(|n| n.to_string()).unwrap_or(EMPTY_TOKEN.to_string()),
               p_head_str.as_ref().map(|s| s.as_ref()).unwrap_or(EMPTY_TOKEN))
    }
}
<|endoftext|>