<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;
extern crate pnet_sys;
extern crate pnet_packet;

use pnet_packet::Packet;
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::ipv4::Ipv4Packet;
use pnet_packet::udp::UdpPacket;
use pnet_packet::icmp::IcmpPacket;
use pnet_packet::icmpv6::Icmpv6Packet;
use pnet_packet::tcp::TcpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};

use std::io;
use std::io::Error;
use std::io::ErrorKind;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;
use std::time::Duration;

/// Represents a transport layer protocol.
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present.
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets.
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol.
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with `transport_channel()`.
pub struct TransportSender {
    pub socket: Arc<pnet_sys::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with `transport_channel()`.
pub struct TransportReceiver {
    pub socket: Arc<pnet_sys::FileDesc>,
    pub buffer: Vec<u8>,
    pub channel_type: TransportChannelType,
}

/// Structure used for holding all configurable options for describing possible options
/// for transport channels.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
        time_to_live: u8,
}

/// Create a new `(TransportSender, TransportReceiver)` pair.
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas `Layer3(IpNextHeaderProtocols::Udp)`
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                pnet_sys::socket(pnet_sys::AF_INET, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                pnet_sys::socket(pnet_sys::AF_INET6, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == pnet_sys::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            pnet_sys::setsockopt(
                socket,
                pnet_sys::IPPROTO_IP,
                pnet_sys::IP_HDRINCL,
                (&hincl as *const libc::c_int) as pnet_sys::Buf,
                mem::size_of::<libc::c_int>() as pnet_sys::SockLen
            )
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(pnet_sys::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: vec![0; buffer_size],
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

/// Create a new `(TransportSender, TransportReceiver)` pair using the additional
/// options specified.
///
/// For a more exhaustive descriptive, see above.
pub fn transport_channel_with(buffer_size: usize,
                              channel_type: TransportChannelType,
                              configuration: Config)
    -> io::Result<(TransportSender, TransportReceiver)> {

    let (sender, receiver) = transport_channel(buffer_size, channel_type)?;

    set_socket_ttl(sender.socket.clone(), configuration.time_to_live)?;
    Ok((sender, receiver))
}

/// Sets a time-to-live for all IP packets sent on the specified socket.
fn set_socket_ttl(socket: Arc<pnet_sys::FileDesc>, ttl: u8) -> io::Result<()> {
    let ttl = ttl as i32;
    let res = unsafe {
        pnet_sys::setsockopt(
            socket.fd,
            pnet_sys::IPPROTO_IP,
            pnet_sys::IP_TTL,
            (&ttl as *const libc::c_int) as pnet_sys::Buf,
            mem::size_of::<libc::c_int>() as pnet_sys::SockLen
        )
    };

    match res {
        -1 => {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket.fd);
            }
            Err(err)
        },
        _ => Ok(()),
    }
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = pnet_sys::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const pnet_sys::SockAddrStorage) as *const pnet_sys::SockAddr;

        pnet_sys::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination.
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    /// Sets a time-to-live on the socket, which then applies for all packets sent.
    pub fn set_ttl(&mut self, time_to_live: u8) -> io::Result<()> {
        set_socket_ttl(self.socket.clone(), time_to_live)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use pnet_packet::MutablePacket;
        use pnet_packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = vec![0; packet.packet().len()];

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```ignore
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty.
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver.
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }

        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel.
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: pnet_sys::SockAddrStorage = unsafe { mem::zeroed() };
                let res = pnet_sys::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = pnet_sys::sockaddr_to_addr(
                            &caddr,
                            mem::size_of::<pnet_sys::SockAddrStorage>()
                        );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use pnet_packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }

            /// Wait only for a timespan of `t` to receive some data, then return. If no data was
            /// received, then `Ok(None)` is returned.
            #[cfg(unix)]
            pub fn next_with_timeout(&mut self, t: Duration) -> io::Result<Option<($ty, IpAddr)>> {
                let socket_fd = self.tr.socket.fd;

                let old_timeout = match pnet_sys::get_socket_receive_timeout(socket_fd) {
                    Err(e) => {
                        eprintln!("Can not get socket timeout before receiving: {}", e);
                        return Err(e)
                    }
                    Ok(t) => t
                };

                match pnet_sys::set_socket_receive_timeout(socket_fd, t) {
                    Err(e) => {
                        eprintln!("Can not set socket timeout for receiving: {}", e);
                        return Err(e)
                    }
                    Ok(_) => {}
                }

                let r = match self.next() {
                    Ok(r) => Ok(Some(r)),
                    Err(e) => match e.kind() {
                        ErrorKind::WouldBlock => Ok(None),
                        _ => {
                            Err(e)
                        }
                    }
                };

                match pnet_sys::set_socket_receive_timeout(socket_fd, old_timeout) {
                    Err(e) => {
                        eprintln!("Can not reset socket timeout after receiving: {}", e);
                    },
                    _ => {}
                };

                r
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);

transport_channel_iterator!(Icmpv6Packet, Icmpv6TransportChannelIterator, icmpv6_packet_iter);

transport_channel_iterator!(TcpPacket, TcpTransportChannelIterator, tcp_packet_iter);
<commit_msg>Properly integrate macro variables in doc strings<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;
extern crate pnet_sys;
extern crate pnet_packet;

use pnet_packet::Packet;
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::ipv4::Ipv4Packet;
use pnet_packet::udp::UdpPacket;
use pnet_packet::icmp::IcmpPacket;
use pnet_packet::icmpv6::Icmpv6Packet;
use pnet_packet::tcp::TcpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};

use std::io;
use std::io::Error;
use std::io::ErrorKind;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;
use std::time::Duration;

/// Represents a transport layer protocol.
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present.
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets.
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol.
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with `transport_channel()`.
pub struct TransportSender {
    pub socket: Arc<pnet_sys::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with `transport_channel()`.
pub struct TransportReceiver {
    pub socket: Arc<pnet_sys::FileDesc>,
    pub buffer: Vec<u8>,
    pub channel_type: TransportChannelType,
}

/// Structure used for holding all configurable options for describing possible options
/// for transport channels.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
        time_to_live: u8,
}

/// Create a new `(TransportSender, TransportReceiver)` pair.
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas `Layer3(IpNextHeaderProtocols::Udp)`
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                pnet_sys::socket(pnet_sys::AF_INET, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                pnet_sys::socket(pnet_sys::AF_INET6, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == pnet_sys::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            pnet_sys::setsockopt(
                socket,
                pnet_sys::IPPROTO_IP,
                pnet_sys::IP_HDRINCL,
                (&hincl as *const libc::c_int) as pnet_sys::Buf,
                mem::size_of::<libc::c_int>() as pnet_sys::SockLen
            )
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(pnet_sys::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: vec![0; buffer_size],
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

/// Create a new `(TransportSender, TransportReceiver)` pair using the additional
/// options specified.
///
/// For a more exhaustive descriptive, see above.
pub fn transport_channel_with(buffer_size: usize,
                              channel_type: TransportChannelType,
                              configuration: Config)
    -> io::Result<(TransportSender, TransportReceiver)> {

    let (sender, receiver) = transport_channel(buffer_size, channel_type)?;

    set_socket_ttl(sender.socket.clone(), configuration.time_to_live)?;
    Ok((sender, receiver))
}

/// Sets a time-to-live for all IP packets sent on the specified socket.
fn set_socket_ttl(socket: Arc<pnet_sys::FileDesc>, ttl: u8) -> io::Result<()> {
    let ttl = ttl as i32;
    let res = unsafe {
        pnet_sys::setsockopt(
            socket.fd,
            pnet_sys::IPPROTO_IP,
            pnet_sys::IP_TTL,
            (&ttl as *const libc::c_int) as pnet_sys::Buf,
            mem::size_of::<libc::c_int>() as pnet_sys::SockLen
        )
    };

    match res {
        -1 => {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket.fd);
            }
            Err(err)
        },
        _ => Ok(()),
    }
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = pnet_sys::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const pnet_sys::SockAddrStorage) as *const pnet_sys::SockAddr;

        pnet_sys::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination.
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    /// Sets a time-to-live on the socket, which then applies for all packets sent.
    pub fn set_ttl(&mut self, time_to_live: u8) -> io::Result<()> {
        set_socket_ttl(self.socket.clone(), time_to_live)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use pnet_packet::MutablePacket;
        use pnet_packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = vec![0; packet.packet().len()];

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```ignore
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        transport_channel_iterator!($ty, $iter, $func, stringify!($ty));
    );
    ($ty:ident, $iter:ident, $func:ident, $tyname:expr) => (
        #[doc = "An iterator over packets of type `"]
        #[doc = $tyname]
        #[doc = "`."]
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }

        #[doc = "Return a packet iterator with packets of type `"]
        #[doc = $tyname]
        #[doc = "` for some transport receiver."]
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }

        impl<'a> $iter<'a> {
            #[doc = "Get the next (`"]
            #[doc = $tyname ]
            #[doc = "`, `IpAddr`) pair for the given channel."]
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: pnet_sys::SockAddrStorage = unsafe { mem::zeroed() };
                let res = pnet_sys::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = pnet_sys::sockaddr_to_addr(
                            &caddr,
                            mem::size_of::<pnet_sys::SockAddrStorage>()
                        );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use pnet_packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }

            /// Wait only for a timespan of `t` to receive some data, then return. If no data was
            /// received, then `Ok(None)` is returned.
            #[cfg(unix)]
            pub fn next_with_timeout(&mut self, t: Duration) -> io::Result<Option<($ty, IpAddr)>> {
                let socket_fd = self.tr.socket.fd;

                let old_timeout = match pnet_sys::get_socket_receive_timeout(socket_fd) {
                    Err(e) => {
                        eprintln!("Can not get socket timeout before receiving: {}", e);
                        return Err(e)
                    }
                    Ok(t) => t
                };

                match pnet_sys::set_socket_receive_timeout(socket_fd, t) {
                    Err(e) => {
                        eprintln!("Can not set socket timeout for receiving: {}", e);
                        return Err(e)
                    }
                    Ok(_) => {}
                }

                let r = match self.next() {
                    Ok(r) => Ok(Some(r)),
                    Err(e) => match e.kind() {
                        ErrorKind::WouldBlock => Ok(None),
                        _ => {
                            Err(e)
                        }
                    }
                };

                match pnet_sys::set_socket_receive_timeout(socket_fd, old_timeout) {
                    Err(e) => {
                        eprintln!("Can not reset socket timeout after receiving: {}", e);
                    },
                    _ => {}
                };

                r
            }
        }
    );
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);

transport_channel_iterator!(Icmpv6Packet, Icmpv6TransportChannelIterator, icmpv6_packet_iter);

transport_channel_iterator!(TcpPacket, TcpTransportChannelIterator, tcp_packet_iter);
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;
extern crate pnet_sys;
extern crate pnet_packet;

use pnet_packet::Packet;
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::ipv4::Ipv4Packet;
use pnet_packet::udp::UdpPacket;
use pnet_packet::icmp::IcmpPacket;
use pnet_packet::icmpv6::Icmpv6Packet;
use pnet_packet::tcp::TcpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<pnet_sys::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<pnet_sys::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                pnet_sys::socket(pnet_sys::AF_INET, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                pnet_sys::socket(pnet_sys::AF_INET6, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == pnet_sys::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            pnet_sys::setsockopt(
                socket,
                pnet_sys::IPPROTO_IP,
                pnet_sys::IP_HDRINCL,
                (&hincl as *const libc::c_int) as pnet_sys::Buf,
                mem::size_of::<libc::c_int>() as pnet_sys::SockLen
            )
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(pnet_sys::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = pnet_sys::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const pnet_sys::SockAddrStorage) as *const pnet_sys::SockAddr;

        pnet_sys::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```ignore
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: pnet_sys::SockAddrStorage = unsafe { mem::zeroed() };
                let res = pnet_sys::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = pnet_sys::sockaddr_to_addr(
                            &caddr,
                            mem::size_of::<pnet_sys::SockAddrStorage>()
                        );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);

transport_channel_iterator!(Icmpv6Packet, Icmpv6TransportChannelIterator, icmpv6_packet_iter);

transport_channel_iterator!(TcpPacket, TcpTransportChannelIterator, tcp_packet_iter);
<commit_msg>Fix macOS build<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;
extern crate pnet_sys;
extern crate pnet_packet;

use pnet_packet::Packet;
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::ipv4::Ipv4Packet;
use pnet_packet::udp::UdpPacket;
use pnet_packet::icmp::IcmpPacket;
use pnet_packet::icmpv6::Icmpv6Packet;
use pnet_packet::tcp::TcpPacket;
use self::TransportChannelType::{Layer3, Layer4};
use self::TransportProtocol::{Ipv4, Ipv6};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::mem;
use std::net::{self, IpAddr};
use std::sync::Arc;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<pnet_sys::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<pnet_sys::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) |
            Layer3(IpNextHeaderProtocol(proto)) => {
                pnet_sys::socket(pnet_sys::AF_INET, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) => {
                pnet_sys::socket(pnet_sys::AF_INET6, pnet_sys::SOCK_RAW, proto as libc::c_int)
            }
        }
    };
    if socket == pnet_sys::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            pnet_sys::setsockopt(
                socket,
                pnet_sys::IPPROTO_IP,
                pnet_sys::IP_HDRINCL,
                (&hincl as *const libc::c_int) as pnet_sys::Buf,
                mem::size_of::<libc::c_int>() as pnet_sys::SockLen
            )
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                pnet_sys::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(pnet_sys::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = pnet_sys::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const pnet_sys::SockAddrStorage) as *const pnet_sys::SockAddr;

        pnet_sys::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use pnet_packet::MutablePacket;
        use pnet_packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```ignore
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: pnet_sys::SockAddrStorage = unsafe { mem::zeroed() };
                let res = pnet_sys::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = pnet_sys::sockaddr_to_addr(
                            &caddr,
                            mem::size_of::<pnet_sys::SockAddrStorage>()
                        );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use pnet_packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet, Ipv4TransportChannelIterator, ipv4_packet_iter);

transport_channel_iterator!(UdpPacket, UdpTransportChannelIterator, udp_packet_iter);

transport_channel_iterator!(IcmpPacket, IcmpTransportChannelIterator, icmp_packet_iter);

transport_channel_iterator!(Icmpv6Packet, Icmpv6TransportChannelIterator, icmpv6_packet_iter);

transport_channel_iterator!(TcpPacket, TcpTransportChannelIterator, tcp_packet_iter);
<|endoftext|>


<commit_before>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

use std::{ptr, os};

use spaces::immix_space::block_info::BlockInfo;

use constants::{BLOCK_SIZE, HEAP_SIZE};
use gc_object::GCObjectRef;

/// The `BlockAllocator` is the global resource for blocks for the immix
/// space.
///
/// On initialization it will allocate a memory map of `HEAP_SIZE` and align
/// it to `BLOCK_SIZE`. During normal runtime it will allocate blocks on the
/// fly from this memory map and store returned blocks in a list.
///
/// Blocks from this `BlockAllocator` are always aligned to `BLOCK_SIZE`.
///
/// The list of returned free blocks is a stack. The `BlockAllocator` will
/// first exhaust the returned free blocks and then fall back to allocating
/// new blocks from the memory map. This means it will return recently
/// returned blocks first.
pub struct BlockAllocator {
    /// The memory map of `HEAP_SIZE`.
    mmap: os::MemoryMap,

    /// The pointer to the last allocated block.
    data: *mut u8,

    /// The upper bound of the memory map as a pointer.
    data_bound: *mut u8,

    /// A list of returned (free) blocks.
    free_blocks: Vec<*mut BlockInfo>,
}

impl BlockAllocator {
    /// Create a new `BlockAllocator`.
    ///
    /// This will `panic` if no memory map of size `HEAP_SIZE` can be
    /// allocared.
    pub fn new() -> BlockAllocator {
        let mmap = os::MemoryMap::new(HEAP_SIZE + BLOCK_SIZE,
                                      &[os::MapOption::MapReadable,
                                        os::MapOption::MapWritable]).unwrap();
        let data = unsafe{ mmap.data().offset((BLOCK_SIZE - (mmap.data() as usize) % BLOCK_SIZE) as isize) };
        let data_bound = unsafe{ mmap.data().offset(mmap.len() as isize) };
        debug!("Allocated heap {:p} of size {}, usable range: {:p} - {:p} (size {}, {} blocks)",
                mmap.data(), mmap.len(), data, data_bound,
                (data_bound as usize) - (data  as usize),
                ((data_bound as usize) - (data  as usize)) / BLOCK_SIZE);
        debug_assert!((data as usize) % BLOCK_SIZE == 0,
            "Allocated mmap {:p} is not aligned (offset {})",
            data, (data as usize) % BLOCK_SIZE);
        return BlockAllocator {
            mmap: mmap,
            data: data,
            data_bound: data_bound,
            free_blocks: Vec::with_capacity(HEAP_SIZE / BLOCK_SIZE),
        };
    }

    /// Get a new block aligned to `BLOCK_SIZE`.
    pub fn get_block(&mut self) -> Option<*mut BlockInfo> {
        return self.free_blocks.pop().or_else(|| self.build_next_block());
    }

    /// Return a collection of blocks.
    pub fn return_blocks(&mut self, blocks: Vec<*mut BlockInfo>) {
        self.free_blocks.extend(blocks.into_iter());
    }

    /// Return the total number of possible blocks.
    pub fn total_blocks(&self) -> usize {
        return HEAP_SIZE / BLOCK_SIZE;
    }

    /// Return the number of unallocated blocks.
    pub fn available_blocks(&self) -> usize {
        return (((self.data_bound as usize) - (self.data as usize)) % BLOCK_SIZE)
            + self.free_blocks.len();
    }

    /// Return if an address is within the bounds of the memory map.
    pub fn is_in_space(&self, object: GCObjectRef) -> bool {
        return self.mmap.data() < (object as *mut u8)
            && (object as *mut u8) < self.data_bound;
    }
}

impl BlockAllocator {
    /// Build a new block from the memory map.
    ///
    /// Returns `None` if the memory map is exhausted.
    fn build_next_block(&mut self) -> Option<*mut BlockInfo> {
        let block = unsafe{ self.data.offset(BLOCK_SIZE as isize) };
        if block < self.data_bound {
            self.data = block;
            debug_assert!((block as usize) % BLOCK_SIZE == 0,
                "Allocated block {:p} is not aligned (offset {})",
                block, (block as usize) % BLOCK_SIZE);
            unsafe{ ptr::write(block as *mut BlockInfo, BlockInfo::new()); }
            return Some(block as *mut BlockInfo);
        }
        return None;
    }
}
<commit_msg>Allow deprecated MemoryMap in block_allocator<commit_after>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

#![allow(deprecated)]
use std::{ptr, os};

use spaces::immix_space::block_info::BlockInfo;

use constants::{BLOCK_SIZE, HEAP_SIZE};
use gc_object::GCObjectRef;

/// The `BlockAllocator` is the global resource for blocks for the immix
/// space.
///
/// On initialization it will allocate a memory map of `HEAP_SIZE` and align
/// it to `BLOCK_SIZE`. During normal runtime it will allocate blocks on the
/// fly from this memory map and store returned blocks in a list.
///
/// Blocks from this `BlockAllocator` are always aligned to `BLOCK_SIZE`.
///
/// The list of returned free blocks is a stack. The `BlockAllocator` will
/// first exhaust the returned free blocks and then fall back to allocating
/// new blocks from the memory map. This means it will return recently
/// returned blocks first.
pub struct BlockAllocator {
    /// The memory map of `HEAP_SIZE`.
    mmap: os::MemoryMap,

    /// The pointer to the last allocated block.
    data: *mut u8,

    /// The upper bound of the memory map as a pointer.
    data_bound: *mut u8,

    /// A list of returned (free) blocks.
    free_blocks: Vec<*mut BlockInfo>,
}

impl BlockAllocator {
    /// Create a new `BlockAllocator`.
    ///
    /// This will `panic` if no memory map of size `HEAP_SIZE` can be
    /// allocared.
    pub fn new() -> BlockAllocator {
        let mmap = os::MemoryMap::new(HEAP_SIZE + BLOCK_SIZE,
                                      &[os::MapOption::MapReadable,
                                        os::MapOption::MapWritable]).unwrap();
        let data = unsafe{ mmap.data().offset((BLOCK_SIZE - (mmap.data() as usize) % BLOCK_SIZE) as isize) };
        let data_bound = unsafe{ mmap.data().offset(mmap.len() as isize) };
        debug!("Allocated heap {:p} of size {}, usable range: {:p} - {:p} (size {}, {} blocks)",
                mmap.data(), mmap.len(), data, data_bound,
                (data_bound as usize) - (data  as usize),
                ((data_bound as usize) - (data  as usize)) / BLOCK_SIZE);
        debug_assert!((data as usize) % BLOCK_SIZE == 0,
            "Allocated mmap {:p} is not aligned (offset {})",
            data, (data as usize) % BLOCK_SIZE);
        return BlockAllocator {
            mmap: mmap,
            data: data,
            data_bound: data_bound,
            free_blocks: Vec::with_capacity(HEAP_SIZE / BLOCK_SIZE),
        };
    }

    /// Get a new block aligned to `BLOCK_SIZE`.
    pub fn get_block(&mut self) -> Option<*mut BlockInfo> {
        return self.free_blocks.pop().or_else(|| self.build_next_block());
    }

    /// Return a collection of blocks.
    pub fn return_blocks(&mut self, blocks: Vec<*mut BlockInfo>) {
        self.free_blocks.extend(blocks.into_iter());
    }

    /// Return the total number of possible blocks.
    pub fn total_blocks(&self) -> usize {
        return HEAP_SIZE / BLOCK_SIZE;
    }

    /// Return the number of unallocated blocks.
    pub fn available_blocks(&self) -> usize {
        return (((self.data_bound as usize) - (self.data as usize)) % BLOCK_SIZE)
            + self.free_blocks.len();
    }

    /// Return if an address is within the bounds of the memory map.
    pub fn is_in_space(&self, object: GCObjectRef) -> bool {
        return self.mmap.data() < (object as *mut u8)
            && (object as *mut u8) < self.data_bound;
    }
}

impl BlockAllocator {
    /// Build a new block from the memory map.
    ///
    /// Returns `None` if the memory map is exhausted.
    fn build_next_block(&mut self) -> Option<*mut BlockInfo> {
        let block = unsafe{ self.data.offset(BLOCK_SIZE as isize) };
        if block < self.data_bound {
            self.data = block;
            debug_assert!((block as usize) % BLOCK_SIZE == 0,
                "Allocated block {:p} is not aligned (offset {})",
                block, (block as usize) % BLOCK_SIZE);
            unsafe{ ptr::write(block as *mut BlockInfo, BlockInfo::new()); }
            return Some(block as *mut BlockInfo);
        }
        return None;
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// compile-flags: -O -C no-prepopulate-passes

#![crate_type="lib"]

struct S;

impl Drop for S {
    fn drop(&mut self) {
    }
}

fn might_unwind() {
}

// CHECK-LABEL: @test
#[no_mangle]
pub fn test() {
    let _s = S;
    // Check that the personality slot alloca gets a lifetime start in each cleanup block, not just
    // in the first one.
    // CHECK-LABEL: cleanup:
    // CHECK: bitcast{{.*}}personalityslot
    // CHECK-NEXT: call void @llvm.lifetime.start
    // CHECK-LABEL: cleanup1:
    // CHECK: bitcast{{.*}}personalityslot
    // CHECK-NEXT: call void @llvm.lifetime.start
    might_unwind();
    might_unwind();
}
<commit_msg>Ignore tests for the personality slot lifetimes on MSVC<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-msvc

// compile-flags: -O -C no-prepopulate-passes

#![crate_type="lib"]

struct S;

impl Drop for S {
    fn drop(&mut self) {
    }
}

fn might_unwind() {
}

// CHECK-LABEL: @test
#[no_mangle]
pub fn test() {
    let _s = S;
    // Check that the personality slot alloca gets a lifetime start in each cleanup block, not just
    // in the first one.
    // CHECK-LABEL: cleanup:
    // CHECK: bitcast{{.*}}personalityslot
    // CHECK-NEXT: call void @llvm.lifetime.start
    // CHECK-LABEL: cleanup1:
    // CHECK: bitcast{{.*}}personalityslot
    // CHECK-NEXT: call void @llvm.lifetime.start
    might_unwind();
    might_unwind();
}
<|endoftext|>


<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Just check we don't get an ICE for the S type.

use std::cell::Cell;
use std::mem;

pub struct S {
    s: Cell<usize>
}

pub const N: usize = 0 - (mem::size_of::<S>() != 4) as usize;
<commit_msg>fix a test comment<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Just check we don't get an ICE for `N`.

use std::cell::Cell;
use std::mem;

pub struct S {
    s: Cell<usize>
}

pub const N: usize = 0 - (mem::size_of::<S>() != 4) as usize;
<|endoftext|>


<commit_before>use crate::ty::SourceType;
use crate::utils::GrowableVec;
use crate::vm::{FileId, NamespaceId, TypeParam};
use dora_parser::ast::{AnnotationParam, Modifier};
use dora_parser::interner::Name;
use dora_parser::Position;
use parking_lot::RwLock;
use std::sync::Arc;

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
pub struct AnnotationId(usize);

impl AnnotationId {
    pub fn max() -> AnnotationId {
        AnnotationId(usize::max_value())
    }
}

impl From<AnnotationId> for usize {
    fn from(data: AnnotationId) -> usize {
        data.0
    }
}

impl From<usize> for AnnotationId {
    fn from(data: usize) -> AnnotationId {
        AnnotationId(data)
    }
}

impl GrowableVec<RwLock<Annotation>> {
    pub fn idx(&self, index: AnnotationId) -> Arc<RwLock<Annotation>> {
        self.idx_usize(index.0)
    }
}

#[derive(Debug)]
pub struct Annotation {
    pub id: AnnotationId,
    pub file_id: FileId,
    pub pos: Position,
    pub name: Name,
    pub namespace_id: NamespaceId,
    pub ty: SourceType,
    pub internal_annotation: Option<Modifier>,

    pub type_params: Option<Vec<TypeParam>>,
    pub term_params: Option<Vec<AnnotationParam>>,
}

impl Annotation {
    pub fn new(
        id: AnnotationId,
        file_id: FileId,
        pos: Position,
        name: Name,
        namespace_id: NamespaceId,
    ) -> Annotation {
        Annotation {
            id,
            file_id,
            pos,
            name,
            namespace_id,
            ty: SourceType::Error,
            internal_annotation: None,
            type_params: None,
            term_params: None,
        }
    }

    pub fn internal(&self) -> bool {
        self.internal_annotation.is_some()
    }
}
<commit_msg>anno: add helper functions to recognize internal annotations<commit_after>use crate::ty::SourceType;
use crate::utils::GrowableVec;
use crate::vm::{FileId, NamespaceId, TypeParam, VM};
use dora_parser::ast::{AnnotationParam, AnnotationUsages, Modifier};
use dora_parser::interner::Name;
use dora_parser::Position;
use parking_lot::RwLock;
use std::sync::Arc;

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
pub struct AnnotationId(usize);

impl AnnotationId {
    pub fn max() -> AnnotationId {
        AnnotationId(usize::max_value())
    }
}

impl From<AnnotationId> for usize {
    fn from(data: AnnotationId) -> usize {
        data.0
    }
}

impl From<usize> for AnnotationId {
    fn from(data: usize) -> AnnotationId {
        AnnotationId(data)
    }
}

impl GrowableVec<RwLock<Annotation>> {
    pub fn idx(&self, index: AnnotationId) -> Arc<RwLock<Annotation>> {
        self.idx_usize(index.0)
    }
}

#[derive(Debug)]
pub struct Annotation {
    pub id: AnnotationId,
    pub file_id: FileId,
    pub pos: Position,
    pub name: Name,
    pub namespace_id: NamespaceId,
    pub ty: SourceType,
    pub internal_annotation: Option<Modifier>,

    pub type_params: Option<Vec<TypeParam>>,
    pub term_params: Option<Vec<AnnotationParam>>,
}

impl Annotation {
    pub fn new(
        id: AnnotationId,
        file_id: FileId,
        pos: Position,
        name: Name,
        namespace_id: NamespaceId,
    ) -> Annotation {
        Annotation {
            id,
            file_id,
            pos,
            name,
            namespace_id,
            ty: SourceType::Error,
            internal_annotation: None,
            type_params: None,
            term_params: None,
        }
    }

    pub fn internal(&self) -> bool {
        self.internal_annotation.is_some()
    }

    pub fn is_abstract(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm
            .annotations
            .idx(vm.known.annotations.abstract_)
            .read()
            .name;
        annotation_usages.contains(name)
    }
    pub fn is_final(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.final_).read().name;
        annotation_usages.contains(name)
    }
    pub fn is_internal(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm
            .annotations
            .idx(vm.known.annotations.internal)
            .read()
            .name;
        annotation_usages.contains(name)
    }
    pub fn is_open(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.open).read().name;
        annotation_usages.contains(name)
    }
    pub fn is_override(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm
            .annotations
            .idx(vm.known.annotations.override_)
            .read()
            .name;
        annotation_usages.contains(name)
    }
    pub fn is_pub(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.pub_).read().name;
        annotation_usages.contains(name)
    }
    pub fn is_static(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.static_).read().name;
        annotation_usages.contains(name)
    }

    pub fn is_test(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.test).read().name;
        annotation_usages.contains(name)
    }

    pub fn is_cannon(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm.annotations.idx(vm.known.annotations.cannon).read().name;
        annotation_usages.contains(name)
    }
    pub fn is_optimize_immediately(annotation_usages: &AnnotationUsages, vm: &VM) -> bool {
        let name = vm
            .annotations
            .idx(vm.known.annotations.optimize_immediately)
            .read()
            .name;
        annotation_usages.contains(name)
    }
}
<|endoftext|>


<commit_before>//! HTTP client
//!
//! HTTPS support is enabled with the `tls` feature.

use super::RequestBuilder;
use crate::{RpcChannel, RpcError, RpcMessage};
use failure::format_err;
use futures::{
	future::{
		self,
		Either::{A, B},
	},
	sync::mpsc,
	Future, Stream,
};
use hyper::{http, rt, Client, Request, Uri};

/// Create a HTTP Client
pub fn connect<TClient>(url: &str) -> impl Future<Item = TClient, Error = RpcError>
where
	TClient: From<RpcChannel>,
{
	let max_parallel = 8;
	let url: Uri = match url.parse() {
		Ok(url) => url,
		Err(e) => return A(future::err(RpcError::Other(e.into()))),
	};

	#[cfg(feature = "tls")]
	let connector = match hyper_tls::HttpsConnector::new(4) {
		Ok(connector) => connector,
		Err(e) => return A(future::err(RpcError::Other(e.into()))),
	};
	#[cfg(feature = "tls")]
	let client = Client::builder().build::<_, hyper::Body>(connector);

	#[cfg(not(feature = "tls"))]
	let client = Client::new();

	let mut request_builder = RequestBuilder::new();

	let (sender, receiver) = mpsc::channel(max_parallel);

	let fut = receiver
		.filter_map(move |msg: RpcMessage| {
			let (request, sender) = match msg {
				RpcMessage::Call(call) => {
					let (_, request) = request_builder.call_request(&call);
					(request, Some(call.sender))
				}
				RpcMessage::Notify(notify) => (request_builder.notification(&notify), None),
				RpcMessage::Subscribe(_) => {
					log::warn!("Unsupported `RpcMessage` type `Subscribe`.");
					return None;
				}
			};

			let request = Request::post(&url)
				.header(
					http::header::CONTENT_TYPE,
					http::header::HeaderValue::from_static("application/json"),
				)
				.body(request.into())
				.expect("Uri and request headers are valid; qed");

			Some(client.request(request).then(move |response| Ok((response, sender))))
		})
		.buffer_unordered(max_parallel)
		.for_each(|(result, sender)| {
			let future = match result {
				Ok(ref res) if !res.status().is_success() => {
					log::trace!("http result status {}", res.status());
					A(future::err(RpcError::Other(format_err!(
						"Unexpected response status code: {}",
						res.status()
					))))
				}
				Ok(res) => B(res
					.into_body()
					.map_err(|e| RpcError::ParseError(e.to_string(), e.into()))
					.concat2()),
				Err(err) => A(future::err(RpcError::Other(err.into()))),
			};
			future.then(|result| {
				if let Some(sender) = sender {
					let response = result
						.and_then(|response| {
							let response_str = String::from_utf8_lossy(response.as_ref()).into_owned();
							super::parse_response(&response_str)
						})
						.and_then(|r| r.1);
					if let Err(err) = sender.send(response) {
						log::warn!("Error resuming asynchronous request: {:?}", err);
					}
				}
				Ok(())
			})
		});

	B(rt::lazy(move || {
		rt::spawn(fut.map_err(|e| log::error!("RPC Client error: {:?}", e)));
		Ok(TClient::from(sender.into()))
	}))
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::*;
	use assert_matches::assert_matches;
	use hyper::rt;
	use jsonrpc_core::{Error, ErrorCode, IoHandler, Params, Value};
	use jsonrpc_http_server::*;
	use std::net::SocketAddr;
	use std::time::Duration;

	fn id<T>(t: T) -> T {
		t
	}

	struct TestServer {
		uri: String,
		socket_addr: SocketAddr,
		server: Option<Server>,
	}

	impl TestServer {
		fn serve<F: FnOnce(ServerBuilder) -> ServerBuilder>(alter: F) -> Self {
			let builder = ServerBuilder::new(io()).rest_api(RestApi::Unsecure);

			let server = alter(builder).start_http(&"127.0.0.1:0".parse().unwrap()).unwrap();
			let socket_addr = server.address().clone();
			let uri = format!("http://{}", socket_addr);

			TestServer {
				uri,
				socket_addr,
				server: Some(server),
			}
		}

		fn start(&mut self) {
			if self.server.is_none() {
				let server = ServerBuilder::new(io())
					.rest_api(RestApi::Unsecure)
					.start_http(&self.socket_addr)
					.unwrap();
				self.server = Some(server);
			} else {
				panic!("Server already running")
			}
		}

		fn stop(&mut self) {
			let server = self.server.take();
			if let Some(server) = server {
				server.close();
			}
		}
	}

	fn io() -> IoHandler {
		let mut io = IoHandler::default();
		io.add_method("hello", |params: Params| match params.parse::<(String,)>() {
			Ok((msg,)) => Ok(Value::String(format!("hello {}", msg))),
			_ => Ok(Value::String("world".into())),
		});
		io.add_method("fail", |_: Params| Err(Error::new(ErrorCode::ServerError(-34))));
		io.add_notification("notify", |params: Params| {
			let (value,) = params.parse::<(u64,)>().expect("expected one u64 as param");
			assert_eq!(value, 12);
		});

		io
	}

	#[derive(Clone)]
	struct TestClient(TypedClient);

	impl From<RpcChannel> for TestClient {
		fn from(channel: RpcChannel) -> Self {
			TestClient(channel.into())
		}
	}

	impl TestClient {
		fn hello(&self, msg: &'static str) -> impl Future<Item = String, Error = RpcError> {
			self.0.call_method("hello", "String", (msg,))
		}
		fn fail(&self) -> impl Future<Item = (), Error = RpcError> {
			self.0.call_method("fail", "()", ())
		}
		fn notify(&self, value: u64) -> impl Future<Item = (), Error = RpcError> {
			self.0.notify("notify", (value,))
		}
	}

	#[test]
	fn should_work() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.hello("http").and_then(move |result| {
					drop(client);
					let _ = tx.send(result);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(run);

		// then
		let result = rx.recv_timeout(Duration::from_secs(3)).unwrap();
		assert_eq!("hello http", result);
	}

	#[test]
	fn should_send_notification() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.notify(12).and_then(move |result| {
					drop(client);
					let _ = tx.send(result);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(run);

		// then
		rx.recv_timeout(Duration::from_secs(3)).unwrap();
	}

	#[test]
	fn handles_invalid_uri() {
		crate::logger::init_log();

		// given
		let invalid_uri = "invalid uri";

		// when
		let run = connect(invalid_uri); // rx.recv_timeout(Duration::from_secs(3)).unwrap();
		let res: Result<TestClient, RpcError> = run.wait();

		// then
		assert_matches!(
			res.map(|_cli| unreachable!()), Err(RpcError::Other(err)) => {
				assert_eq!(format!("{}", err), "invalid uri character");
			}
		);
	}

	#[test]
	fn handles_server_error() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.fail().then(move |res| {
					let _ = tx.send(res);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));
		rt::run(run);

		// then
		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();

		if let Err(RpcError::JsonRpcError(err)) = res {
			assert_eq!(
				err,
				Error {
					code: ErrorCode::ServerError(-34),
					message: "Server error".into(),
					data: None
				}
			)
		} else {
			panic!("Expected JsonRpcError. Received {:?}", res)
		}
	}

	#[test]
	fn handles_connection_refused_error() {
		// given
		let mut server = TestServer::serve(id);
		// stop server so that we get a connection refused
		server.stop();
		let (tx, rx) = std::sync::mpsc::channel();

		let client = connect(&server.uri);

		let call = client
			.and_then(|client: TestClient| {
				client.hello("http").then(move |res| {
					let _ = tx.send(res);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(call);

		// then
		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();

		if let Err(RpcError::Other(err)) = res {
			if let Some(err) = err.downcast_ref::<hyper::Error>() {
				assert!(err.is_connect(), format!("Expected Connection Error, got {:?}", err))
			} else {
				panic!("Expected a hyper::Error")
			}
		} else {
			panic!("Expected JsonRpcError. Received {:?}", res)
		}
	}

	#[test]
	#[ignore] // todo: [AJ] make it pass
	fn client_still_works_after_http_connect_error() {
		// given
		let mut server = TestServer::serve(id);

		// stop server so that we get a connection refused
		server.stop();

		let (tx, rx) = std::sync::mpsc::channel();
		let tx2 = tx.clone();

		let client = connect(&server.uri);

		let call = client
			.and_then(move |client: TestClient| {
				client
					.hello("http")
					.then(move |res| {
						let _ = tx.send(res);
						Ok(())
					})
					.and_then(move |_| {
						server.start(); // todo: make the server start on the main thread
						client.hello("http2").then(move |res| {
							let _ = tx2.send(res);
							Ok(())
						})
					})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		// when
		rt::run(call);

		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();
		assert!(res.is_err());

		// then
		let result = rx.recv_timeout(Duration::from_secs(3)).unwrap().unwrap();
		assert_eq!("hello http", result);
	}
}
<commit_msg>Add "Accept:application/json" header for all requests. (#544)<commit_after>//! HTTP client
//!
//! HTTPS support is enabled with the `tls` feature.

use super::RequestBuilder;
use crate::{RpcChannel, RpcError, RpcMessage};
use failure::format_err;
use futures::{
	future::{
		self,
		Either::{A, B},
	},
	sync::mpsc,
	Future, Stream,
};
use hyper::{http, rt, Client, Request, Uri};

/// Create a HTTP Client
pub fn connect<TClient>(url: &str) -> impl Future<Item = TClient, Error = RpcError>
where
	TClient: From<RpcChannel>,
{
	let max_parallel = 8;
	let url: Uri = match url.parse() {
		Ok(url) => url,
		Err(e) => return A(future::err(RpcError::Other(e.into()))),
	};

	#[cfg(feature = "tls")]
	let connector = match hyper_tls::HttpsConnector::new(4) {
		Ok(connector) => connector,
		Err(e) => return A(future::err(RpcError::Other(e.into()))),
	};
	#[cfg(feature = "tls")]
	let client = Client::builder().build::<_, hyper::Body>(connector);

	#[cfg(not(feature = "tls"))]
	let client = Client::new();

	let mut request_builder = RequestBuilder::new();

	let (sender, receiver) = mpsc::channel(max_parallel);

	let fut = receiver
		.filter_map(move |msg: RpcMessage| {
			let (request, sender) = match msg {
				RpcMessage::Call(call) => {
					let (_, request) = request_builder.call_request(&call);
					(request, Some(call.sender))
				}
				RpcMessage::Notify(notify) => (request_builder.notification(&notify), None),
				RpcMessage::Subscribe(_) => {
					log::warn!("Unsupported `RpcMessage` type `Subscribe`.");
					return None;
				}
			};

			let request = Request::post(&url)
				.header(
					http::header::CONTENT_TYPE,
					http::header::HeaderValue::from_static("application/json"),
				)
				.header(
					http::header::ACCEPT,
					http::header::HeaderValue::from_static("application/json"),
				)
				.body(request.into())
				.expect("Uri and request headers are valid; qed");

			Some(client.request(request).then(move |response| Ok((response, sender))))
		})
		.buffer_unordered(max_parallel)
		.for_each(|(result, sender)| {
			let future = match result {
				Ok(ref res) if !res.status().is_success() => {
					log::trace!("http result status {}", res.status());
					A(future::err(RpcError::Other(format_err!(
						"Unexpected response status code: {}",
						res.status()
					))))
				}
				Ok(res) => B(res
					.into_body()
					.map_err(|e| RpcError::ParseError(e.to_string(), e.into()))
					.concat2()),
				Err(err) => A(future::err(RpcError::Other(err.into()))),
			};
			future.then(|result| {
				if let Some(sender) = sender {
					let response = result
						.and_then(|response| {
							let response_str = String::from_utf8_lossy(response.as_ref()).into_owned();
							super::parse_response(&response_str)
						})
						.and_then(|r| r.1);
					if let Err(err) = sender.send(response) {
						log::warn!("Error resuming asynchronous request: {:?}", err);
					}
				}
				Ok(())
			})
		});

	B(rt::lazy(move || {
		rt::spawn(fut.map_err(|e| log::error!("RPC Client error: {:?}", e)));
		Ok(TClient::from(sender.into()))
	}))
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::*;
	use assert_matches::assert_matches;
	use hyper::rt;
	use jsonrpc_core::{Error, ErrorCode, IoHandler, Params, Value};
	use jsonrpc_http_server::*;
	use std::net::SocketAddr;
	use std::time::Duration;

	fn id<T>(t: T) -> T {
		t
	}

	struct TestServer {
		uri: String,
		socket_addr: SocketAddr,
		server: Option<Server>,
	}

	impl TestServer {
		fn serve<F: FnOnce(ServerBuilder) -> ServerBuilder>(alter: F) -> Self {
			let builder = ServerBuilder::new(io()).rest_api(RestApi::Unsecure);

			let server = alter(builder).start_http(&"127.0.0.1:0".parse().unwrap()).unwrap();
			let socket_addr = server.address().clone();
			let uri = format!("http://{}", socket_addr);

			TestServer {
				uri,
				socket_addr,
				server: Some(server),
			}
		}

		fn start(&mut self) {
			if self.server.is_none() {
				let server = ServerBuilder::new(io())
					.rest_api(RestApi::Unsecure)
					.start_http(&self.socket_addr)
					.unwrap();
				self.server = Some(server);
			} else {
				panic!("Server already running")
			}
		}

		fn stop(&mut self) {
			let server = self.server.take();
			if let Some(server) = server {
				server.close();
			}
		}
	}

	fn io() -> IoHandler {
		let mut io = IoHandler::default();
		io.add_method("hello", |params: Params| match params.parse::<(String,)>() {
			Ok((msg,)) => Ok(Value::String(format!("hello {}", msg))),
			_ => Ok(Value::String("world".into())),
		});
		io.add_method("fail", |_: Params| Err(Error::new(ErrorCode::ServerError(-34))));
		io.add_notification("notify", |params: Params| {
			let (value,) = params.parse::<(u64,)>().expect("expected one u64 as param");
			assert_eq!(value, 12);
		});

		io
	}

	#[derive(Clone)]
	struct TestClient(TypedClient);

	impl From<RpcChannel> for TestClient {
		fn from(channel: RpcChannel) -> Self {
			TestClient(channel.into())
		}
	}

	impl TestClient {
		fn hello(&self, msg: &'static str) -> impl Future<Item = String, Error = RpcError> {
			self.0.call_method("hello", "String", (msg,))
		}
		fn fail(&self) -> impl Future<Item = (), Error = RpcError> {
			self.0.call_method("fail", "()", ())
		}
		fn notify(&self, value: u64) -> impl Future<Item = (), Error = RpcError> {
			self.0.notify("notify", (value,))
		}
	}

	#[test]
	fn should_work() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.hello("http").and_then(move |result| {
					drop(client);
					let _ = tx.send(result);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(run);

		// then
		let result = rx.recv_timeout(Duration::from_secs(3)).unwrap();
		assert_eq!("hello http", result);
	}

	#[test]
	fn should_send_notification() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.notify(12).and_then(move |result| {
					drop(client);
					let _ = tx.send(result);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(run);

		// then
		rx.recv_timeout(Duration::from_secs(3)).unwrap();
	}

	#[test]
	fn handles_invalid_uri() {
		crate::logger::init_log();

		// given
		let invalid_uri = "invalid uri";

		// when
		let run = connect(invalid_uri); // rx.recv_timeout(Duration::from_secs(3)).unwrap();
		let res: Result<TestClient, RpcError> = run.wait();

		// then
		assert_matches!(
			res.map(|_cli| unreachable!()), Err(RpcError::Other(err)) => {
				assert_eq!(format!("{}", err), "invalid uri character");
			}
		);
	}

	#[test]
	fn handles_server_error() {
		crate::logger::init_log();

		// given
		let server = TestServer::serve(id);
		let (tx, rx) = std::sync::mpsc::channel();

		// when
		let run = connect(&server.uri)
			.and_then(|client: TestClient| {
				client.fail().then(move |res| {
					let _ = tx.send(res);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));
		rt::run(run);

		// then
		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();

		if let Err(RpcError::JsonRpcError(err)) = res {
			assert_eq!(
				err,
				Error {
					code: ErrorCode::ServerError(-34),
					message: "Server error".into(),
					data: None
				}
			)
		} else {
			panic!("Expected JsonRpcError. Received {:?}", res)
		}
	}

	#[test]
	fn handles_connection_refused_error() {
		// given
		let mut server = TestServer::serve(id);
		// stop server so that we get a connection refused
		server.stop();
		let (tx, rx) = std::sync::mpsc::channel();

		let client = connect(&server.uri);

		let call = client
			.and_then(|client: TestClient| {
				client.hello("http").then(move |res| {
					let _ = tx.send(res);
					Ok(())
				})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		rt::run(call);

		// then
		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();

		if let Err(RpcError::Other(err)) = res {
			if let Some(err) = err.downcast_ref::<hyper::Error>() {
				assert!(err.is_connect(), format!("Expected Connection Error, got {:?}", err))
			} else {
				panic!("Expected a hyper::Error")
			}
		} else {
			panic!("Expected JsonRpcError. Received {:?}", res)
		}
	}

	#[test]
	#[ignore] // todo: [AJ] make it pass
	fn client_still_works_after_http_connect_error() {
		// given
		let mut server = TestServer::serve(id);

		// stop server so that we get a connection refused
		server.stop();

		let (tx, rx) = std::sync::mpsc::channel();
		let tx2 = tx.clone();

		let client = connect(&server.uri);

		let call = client
			.and_then(move |client: TestClient| {
				client
					.hello("http")
					.then(move |res| {
						let _ = tx.send(res);
						Ok(())
					})
					.and_then(move |_| {
						server.start(); // todo: make the server start on the main thread
						client.hello("http2").then(move |res| {
							let _ = tx2.send(res);
							Ok(())
						})
					})
			})
			.map_err(|e| log::error!("RPC Client error: {:?}", e));

		// when
		rt::run(call);

		let res = rx.recv_timeout(Duration::from_secs(3)).unwrap();
		assert!(res.is_err());

		// then
		let result = rx.recv_timeout(Duration::from_secs(3)).unwrap().unwrap();
		assert_eq!("hello http", result);
	}
}
<|endoftext|>