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

/// This example shows a basic packet logger using libpnet

extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::Packet;
use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use std::env;
use std::io::{self, Write};
use std::process;
use std::net::IpAddr;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!("[{}]: UDP Packet: {}:{} > {}:{}; length: {} payload: {}",
                 interface_name,
                 source,
                 udp.get_source(),
                 destination,
                 udp.get_destination(),
                 udp.get_length(),
                 udp.get_payload());
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_reply_packet.get_sequence_number(),
                         echo_reply_packet.get_identifier());
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_request_packet.get_sequence_number(),
                         echo_request_packet.get_identifier());
            }
            _ => {
                println!("[{}]: ICMP packet {} -> {} (type={:?})",
                         interface_name,
                         source,
                         destination,
                         icmp_packet.get_icmp_type())
            }
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!("[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 tcp.get_source(),
                 destination,
                 tcp.get_destination(),
                 packet.len());
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(interface_name: &str,
                             source: IpAddr,
                             destination: IpAddr,
                             protocol: IpNextHeaderProtocol,
                             packet: &[u8]) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => {
            println!("[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
                     interface_name,
                     match source {
                         IpAddr::V4(..) => "IPv4",
                         _ => "IPv6",
                     },
                     source,
                     destination,
                     protocol,
                     packet.len())
        }

    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V4(header.get_source()),
                                  IpAddr::V4(header.get_destination()),
                                  header.get_next_level_protocol(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V6(header.get_source()),
                                  IpAddr::V6(header.get_destination()),
                                  header.get_next_header(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!("[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
                 interface_name,
                 ethernet.get_source(),
                 header.get_sender_proto_addr(),
                 ethernet.get_destination(),
                 header.get_target_proto_addr(),
                 header.get_operation());
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_packet(interface_name: &str, ethernet: &EthernetPacket) {
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => {
            println!("[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
                     interface_name,
                     ethernet.get_source(),
                     ethernet.get_destination(),
                     ethernet.get_ethertype(),
                     ethernet.packet().len())
        }
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        },
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces.into_iter().filter(interface_names_match).next().unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        match rx.next() {
            Ok(packet) => handle_packet(&interface.name[..], &EthernetPacket::new(packet).unwrap()),
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<commit_msg>Update main.rs<commit_after>extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::Packet;
use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use std::env;
use std::io::{self, Write};
use std::process;
use std::net::IpAddr;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!("[{}]: UDP Packet: {}:{} > {}:{}; length: {} payload: {}",
                 interface_name,
                 source,
                 udp.get_source(),
                 destination,
                 udp.get_destination(),
                 udp.get_length(),
                 udp.get_payload());
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_reply_packet.get_sequence_number(),
                         echo_reply_packet.get_identifier());
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_request_packet.get_sequence_number(),
                         echo_request_packet.get_identifier());
            }
            _ => {
                println!("[{}]: ICMP packet {} -> {} (type={:?})",
                         interface_name,
                         source,
                         destination,
                         icmp_packet.get_icmp_type())
            }
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!("[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 tcp.get_source(),
                 destination,
                 tcp.get_destination(),
                 packet.len());
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(interface_name: &str,
                             source: IpAddr,
                             destination: IpAddr,
                             protocol: IpNextHeaderProtocol,
                             packet: &[u8]) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => {
            println!("[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
                     interface_name,
                     match source {
                         IpAddr::V4(..) => "IPv4",
                         _ => "IPv6",
                     },
                     source,
                     destination,
                     protocol,
                     packet.len())
        }

    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V4(header.get_source()),
                                  IpAddr::V4(header.get_destination()),
                                  header.get_next_level_protocol(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V6(header.get_source()),
                                  IpAddr::V6(header.get_destination()),
                                  header.get_next_header(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!("[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
                 interface_name,
                 ethernet.get_source(),
                 header.get_sender_proto_addr(),
                 ethernet.get_destination(),
                 header.get_target_proto_addr(),
                 header.get_operation());
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_packet(interface_name: &str, ethernet: &EthernetPacket) {
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => {
            println!("[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
                     interface_name,
                     ethernet.get_source(),
                     ethernet.get_destination(),
                     ethernet.get_ethertype(),
                     ethernet.packet().len())
        }
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        },
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces.into_iter().filter(interface_names_match).next().unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        match rx.next() {
            Ok(packet) => handle_packet(&interface.name[..], &EthernetPacket::new(packet).unwrap()),
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<|endoftext|>


<commit_before>pub mod edge;
pub mod face;
pub mod path;
mod traverse;
pub mod vertex;

use fool::BoolExt;
use std::ops::{Deref, DerefMut};

use crate::graph::borrow::{Reborrow, ReborrowMut};
use crate::graph::mutation::Consistent;
use crate::graph::storage::key::OpaqueKey;
use crate::graph::storage::payload::Entity;
use crate::graph::storage::{AsStorage, AsStorageMut};
use crate::graph::GraphError;

// TODO: Use `bind_unchecked` whenever possible (that is, when it is logically
//       consistent to assume that the key is present in storage).
// TODO: Consider `Bind` and `Unbind` traits and decomposing the `ClosedView`
//       trait.

pub trait ClosedView: Deref<Target = <Self as ClosedView>::Entity> {
    type Key: OpaqueKey;
    type Entity: Entity<Key = Self::Key>;

    /// Gets the key for the view.
    fn key(&self) -> Self::Key;

    /// Rebinds a view's storage with the given key.
    ///
    /// Rebinding a view allows its underlying storage to be reinterpretted.
    /// The output view must also be bound to a payload.
    ///
    /// # Examples
    ///
    /// Perform a fallible traversal and preserve mutability of the resulting
    /// view:
    ///
    /// ```rust,no_run
    /// # use plexus::graph::{ClosedView, MeshGraph};
    /// # use plexus::prelude::*;
    /// #
    /// # let mut graph = MeshGraph::<()>::default();
    /// # let key = graph.faces().keys().nth(0).unwrap();
    /// // ...
    /// let face = graph.face_mut(key).unwrap();
    /// // Find a face along a boundary. If no such face is found, continue to use the
    /// // initiating face.
    /// let mut face = {
    ///     let key = face
    ///         .traverse_by_depth()
    ///         .find(|face| {
    ///             face.interior_arcs()
    ///                 .map(|arc| arc.into_opposite_arc())
    ///                 .any(|arc| arc.is_boundary_arc())
    ///         })
    ///         .map(|face| face.key());
    ///     if let Some(key) = key {
    ///         face.rebind(key).unwrap() // Rebind into the boundary face.
    ///     }
    ///     else {
    ///         face
    ///     }
    /// };
    /// ```
    fn rebind<T, M>(self, key: T::Key) -> Result<T, GraphError>
    where
        Self: Into<View<M, <Self as ClosedView>::Entity>>,
        T: From<View<M, <T as ClosedView>::Entity>> + ClosedView,
        M: Reborrow,
        M::Target: AsStorage<Self::Entity> + AsStorage<T::Entity>,
    {
        self.into()
            .rebind_into::<_, T::Entity>(key)
            .ok_or_else(|| GraphError::TopologyNotFound)
    }
}

pub struct View<M, T>
where
    M: Reborrow,
    M::Target: AsStorage<T>,
    T: Entity,
{
    storage: M,
    key: T::Key,
}

impl<M, T> View<M, T>
where
    M: Reborrow,
    M::Target: AsStorage<T>,
    T: Entity,
{
    pub fn bind(storage: M, key: T::Key) -> Option<Self> {
        storage
            .reborrow()
            .as_storage()
            .contains_key(&key)
            .some(View::bind_unchecked(storage, key))
    }

    pub fn bind_into<U>(storage: M, key: T::Key) -> Option<U>
    where
        U: From<Self>,
    {
        View::bind(storage, key).map(U::from)
    }

    pub fn unbind(self) -> (M, T::Key) {
        let View { storage, key, .. } = self;
        (storage, key)
    }

    pub fn rebind<U>(self, key: U::Key) -> Option<View<M, U>>
    where
        U: Entity,
        M::Target: AsStorage<U>,
    {
        let (storage, _) = self.unbind();
        View::bind(storage, key)
    }

    pub fn rebind_into<V, U>(self, key: U::Key) -> Option<V>
    where
        V: From<View<M, U>>,
        U: Entity,
        M::Target: AsStorage<U>,
    {
        self.rebind(key).map(V::from)
    }

    pub fn key(&self) -> T::Key {
        self.key
    }

    pub fn interior_reborrow(&self) -> View<&M::Target, T> {
        View::bind_unchecked(self.storage.reborrow(), self.key)
    }

    pub(in crate::graph) fn bind_unchecked(storage: M, key: T::Key) -> Self {
        View { storage, key }
    }
}

impl<M, T> View<M, T>
where
    M: ReborrowMut,
    M::Target: AsStorage<T>,
    T: Entity,
{
    pub fn interior_reborrow_mut(&mut self) -> View<&mut M::Target, T> {
        View::bind_unchecked(self.storage.reborrow_mut(), self.key)
    }
}

impl<'a, M, T> View<&'a mut M, T>
where
    M: 'a + AsStorageMut<T>,
    T: 'a + Entity,
{
    pub fn into_ref(self) -> View<&'a M, T> {
        let (storage, key) = self.unbind();
        View::bind(&*storage, key).unwrap()
    }
}

impl<M, T> Clone for View<M, T>
where
    M: Clone + Reborrow,
    M::Target: AsStorage<T>,
    T: Entity,
{
    fn clone(&self) -> Self {
        View {
            storage: self.storage.clone(),
            key: self.key,
        }
    }
}

impl<M, T> Copy for View<M, T>
where
    M: Copy + Reborrow,
    M::Target: AsStorage<T>,
    T: Entity,
{
}

impl<M, T> Deref for View<M, T>
where
    M: Reborrow,
    M::Target: AsStorage<T>,
    T: Entity,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.storage
            .reborrow()
            .as_storage()
            .get(&self.key)
            .expect("view key invalidated")
    }
}

impl<M, T> DerefMut for View<M, T>
where
    M: ReborrowMut,
    M::Target: AsStorageMut<T>,
    T: Entity,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.storage
            .reborrow_mut()
            .as_storage_mut()
            .get_mut(&self.key)
            .expect("view key invalidated")
    }
}

// TODO: Consider implementing `Eq` for views.
impl<M, T> PartialEq for View<M, T>
where
    M: Reborrow,
    M::Target: AsStorage<T> + Consistent,
    T: Entity,
{
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key
    }
}

pub struct Orphan<'a, T>
where
    T: Entity,
{
    payload: &'a mut T,
    key: T::Key,
}

impl<'a, T> Orphan<'a, T>
where
    T: 'a + Entity,
{
    pub fn bind<M>(storage: &'a mut M, key: T::Key) -> Option<Self>
    where
        M: AsStorageMut<T>,
    {
        View::bind(storage, key).map(Orphan::from)
    }

    pub fn bind_into<U, M>(storage: &'a mut M, key: T::Key) -> Option<U>
    where
        U: From<Self>,
        M: AsStorageMut<T>,
    {
        Orphan::bind(storage, key).map(U::from)
    }

    pub fn key(&self) -> T::Key {
        self.key
    }

    pub(in crate::graph) fn bind_unchecked(payload: &'a mut T, key: T::Key) -> Self {
        Orphan { payload, key }
    }
}

impl<'a, T> Deref for Orphan<'a, T>
where
    T: 'a + Entity,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &*self.payload
    }
}

impl<'a, T> DerefMut for Orphan<'a, T>
where
    T: 'a + Entity,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.payload
    }
}

impl<'a, T, M> From<View<&'a mut M, T>> for Orphan<'a, T>
where
    T: 'a + Entity,
    M: AsStorageMut<T>,
{
    fn from(view: View<&'a mut M, T>) -> Self {
        let (storage, key) = view.unbind();
        let payload = storage
            .as_storage_mut()
            .get_mut(&key)
            .expect("view key invalidated");
        Orphan::bind_unchecked(payload, key)
    }
}
<commit_msg>Rename `View` type parameters.<commit_after>pub mod edge;
pub mod face;
pub mod path;
mod traverse;
pub mod vertex;

use fool::BoolExt;
use std::ops::{Deref, DerefMut};

use crate::graph::borrow::{Reborrow, ReborrowMut};
use crate::graph::mutation::Consistent;
use crate::graph::storage::key::OpaqueKey;
use crate::graph::storage::payload::Entity;
use crate::graph::storage::{AsStorage, AsStorageMut};
use crate::graph::GraphError;

// TODO: Use `bind_unchecked` whenever possible (that is, when it is logically
//       consistent to assume that the key is present in storage).
// TODO: Consider `Bind` and `Unbind` traits and decomposing the `ClosedView`
//       trait.

pub trait ClosedView: Deref<Target = <Self as ClosedView>::Entity> {
    type Key: OpaqueKey;
    type Entity: Entity<Key = Self::Key>;

    /// Gets the key for the view.
    fn key(&self) -> Self::Key;

    /// Rebinds a view's storage with the given key.
    ///
    /// Rebinding a view allows its underlying storage to be reinterpretted.
    /// The output view must also be bound to a payload.
    ///
    /// # Examples
    ///
    /// Perform a fallible traversal and preserve mutability of the resulting
    /// view:
    ///
    /// ```rust,no_run
    /// # use plexus::graph::{ClosedView, MeshGraph};
    /// # use plexus::prelude::*;
    /// #
    /// # let mut graph = MeshGraph::<()>::default();
    /// # let key = graph.faces().keys().nth(0).unwrap();
    /// // ...
    /// let face = graph.face_mut(key).unwrap();
    /// // Find a face along a boundary. If no such face is found, continue to use the
    /// // initiating face.
    /// let mut face = {
    ///     let key = face
    ///         .traverse_by_depth()
    ///         .find(|face| {
    ///             face.interior_arcs()
    ///                 .map(|arc| arc.into_opposite_arc())
    ///                 .any(|arc| arc.is_boundary_arc())
    ///         })
    ///         .map(|face| face.key());
    ///     if let Some(key) = key {
    ///         face.rebind(key).unwrap() // Rebind into the boundary face.
    ///     }
    ///     else {
    ///         face
    ///     }
    /// };
    /// ```
    fn rebind<T, B>(self, key: T::Key) -> Result<T, GraphError>
    where
        Self: Into<View<B, <Self as ClosedView>::Entity>>,
        T: From<View<B, <T as ClosedView>::Entity>> + ClosedView,
        B: Reborrow,
        B::Target: AsStorage<Self::Entity> + AsStorage<T::Entity>,
    {
        self.into()
            .rebind_into::<_, T::Entity>(key)
            .ok_or_else(|| GraphError::TopologyNotFound)
    }
}

pub struct View<B, E>
where
    B: Reborrow,
    B::Target: AsStorage<E>,
    E: Entity,
{
    storage: B,
    key: E::Key,
}

impl<B, E> View<B, E>
where
    B: Reborrow,
    B::Target: AsStorage<E>,
    E: Entity,
{
    pub fn bind(storage: B, key: E::Key) -> Option<Self> {
        storage
            .reborrow()
            .as_storage()
            .contains_key(&key)
            .some(View::bind_unchecked(storage, key))
    }

    pub fn bind_into<U>(storage: B, key: E::Key) -> Option<U>
    where
        U: From<Self>,
    {
        View::bind(storage, key).map(U::from)
    }

    pub fn unbind(self) -> (B, E::Key) {
        let View { storage, key, .. } = self;
        (storage, key)
    }

    pub fn rebind<U>(self, key: U::Key) -> Option<View<B, U>>
    where
        U: Entity,
        B::Target: AsStorage<U>,
    {
        let (storage, _) = self.unbind();
        View::bind(storage, key)
    }

    pub fn rebind_into<V, U>(self, key: U::Key) -> Option<V>
    where
        V: From<View<B, U>>,
        U: Entity,
        B::Target: AsStorage<U>,
    {
        self.rebind(key).map(V::from)
    }

    pub fn key(&self) -> E::Key {
        self.key
    }

    pub fn interior_reborrow(&self) -> View<&B::Target, E> {
        View::bind_unchecked(self.storage.reborrow(), self.key)
    }

    pub(in crate::graph) fn bind_unchecked(storage: B, key: E::Key) -> Self {
        View { storage, key }
    }
}

impl<B, E> View<B, E>
where
    B: ReborrowMut,
    B::Target: AsStorage<E>,
    E: Entity,
{
    pub fn interior_reborrow_mut(&mut self) -> View<&mut B::Target, E> {
        View::bind_unchecked(self.storage.reborrow_mut(), self.key)
    }
}

impl<'a, M, E> View<&'a mut M, E>
where
    M: 'a + AsStorageMut<E>,
    E: 'a + Entity,
{
    pub fn into_ref(self) -> View<&'a M, E> {
        let (storage, key) = self.unbind();
        View::bind(&*storage, key).unwrap()
    }
}

impl<B, E> Clone for View<B, E>
where
    B: Clone + Reborrow,
    B::Target: AsStorage<E>,
    E: Entity,
{
    fn clone(&self) -> Self {
        View {
            storage: self.storage.clone(),
            key: self.key,
        }
    }
}

impl<B, E> Copy for View<B, E>
where
    B: Copy + Reborrow,
    B::Target: AsStorage<E>,
    E: Entity,
{
}

impl<B, E> Deref for View<B, E>
where
    B: Reborrow,
    B::Target: AsStorage<E>,
    E: Entity,
{
    type Target = E;

    fn deref(&self) -> &Self::Target {
        self.storage
            .reborrow()
            .as_storage()
            .get(&self.key)
            .expect("view key invalidated")
    }
}

impl<B, E> DerefMut for View<B, E>
where
    B: ReborrowMut,
    B::Target: AsStorageMut<E>,
    E: Entity,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.storage
            .reborrow_mut()
            .as_storage_mut()
            .get_mut(&self.key)
            .expect("view key invalidated")
    }
}

// TODO: Consider implementing `Eq` for views.
impl<B, E> PartialEq for View<B, E>
where
    B: Reborrow,
    B::Target: AsStorage<E> + Consistent,
    E: Entity,
{
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key
    }
}

pub struct Orphan<'a, E>
where
    E: Entity,
{
    payload: &'a mut E,
    key: E::Key,
}

impl<'a, E> Orphan<'a, E>
where
    E: 'a + Entity,
{
    pub fn bind<M>(storage: &'a mut M, key: E::Key) -> Option<Self>
    where
        M: AsStorageMut<E>,
    {
        View::bind(storage, key).map(Orphan::from)
    }

    pub fn bind_into<U, M>(storage: &'a mut M, key: E::Key) -> Option<U>
    where
        U: From<Self>,
        M: AsStorageMut<E>,
    {
        Orphan::bind(storage, key).map(U::from)
    }

    pub fn key(&self) -> E::Key {
        self.key
    }

    pub(in crate::graph) fn bind_unchecked(payload: &'a mut E, key: E::Key) -> Self {
        Orphan { payload, key }
    }
}

impl<'a, E> Deref for Orphan<'a, E>
where
    E: 'a + Entity,
{
    type Target = E;

    fn deref(&self) -> &Self::Target {
        &*self.payload
    }
}

impl<'a, E> DerefMut for Orphan<'a, E>
where
    E: 'a + Entity,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.payload
    }
}

impl<'a, E, M> From<View<&'a mut M, E>> for Orphan<'a, E>
where
    E: 'a + Entity,
    M: AsStorageMut<E>,
{
    fn from(view: View<&'a mut M, E>) -> Self {
        let (storage, key) = view.unbind();
        let payload = storage
            .as_storage_mut()
            .get_mut(&key)
            .expect("view key invalidated");
        Orphan::bind_unchecked(payload, key)
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
//    https://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 autocxx_parser::file_locations::FileLocationStrategy;
use proc_macro2::TokenStream;

use crate::{ParseError, ParsedFile, RebuildDependencyRecorder};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::{ffi::OsStr, io, process};
use std::{fmt::Display, fs::File};

/// Errors returned during creation of a cc::Build from an include_cxx
/// macro.
#[derive(Debug)]
pub enum BuilderError {
    /// The cxx module couldn't parse the code generated by autocxx.
    /// This could well be a bug in autocxx.
    InvalidCxx(cxx_gen::Error),
    /// The .rs file didn't exist or couldn't be parsed.
    ParseError(ParseError),
    /// We couldn't write the c++ code to disk.
    FileWriteFail(std::io::Error, PathBuf),
    /// No `include_cxx` macro was found anywhere.
    NoIncludeCxxMacrosFound,
    /// Unable to create one of the directories to which we need to write
    UnableToCreateDirectory(std::io::Error, PathBuf),
}

impl Display for BuilderError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BuilderError::ParseError(pe) => write!(f, "Unable to parse .rs file: {}", pe)?,
            BuilderError::InvalidCxx(ee) => write!(f, "cxx was unable to understand the code generated by autocxx (likely a bug in autocxx; please report.) {}", ee)?,
            BuilderError::FileWriteFail(ee, pb) => write!(f, "Unable to write to {}: {}", pb.to_string_lossy(), ee)?,
            BuilderError::NoIncludeCxxMacrosFound => write!(f, "No include_cpp! macro found")?,
            BuilderError::UnableToCreateDirectory(ee, pb) => write!(f, "Unable to create directory {}: {}", pb.to_string_lossy(), ee)?,
        }
        Ok(())
    }
}

pub type BuilderBuild = cc::Build;

pub struct BuilderSuccess(pub BuilderBuild, pub Vec<PathBuf>);

/// Results of a build.
pub type BuilderResult = Result<BuilderSuccess, BuilderError>;

/// Build autocxx C++ files and return a cc::Build you can use to build
/// more from a build.rs file.
/// You need to provide the Rust file path and the iterator of paths
/// which should be used as include directories.
pub fn build<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderResult
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    build_to_custom_directory(
        rs_file,
        autocxx_incs,
        extra_clang_args,
        None,
        dependency_recorder,
    )
}

/// Builds successfully, or exits the process displaying a suitable
/// message.
pub fn expect_build<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderSuccess
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    build(rs_file, autocxx_incs, extra_clang_args, dependency_recorder).unwrap_or_else(|err| {
        let _ = writeln!(io::stderr(), "\n\nautocxx error: {}\n\n", report(err));
        process::exit(1);
    })
}

fn report(err: BuilderError) -> String {
    err.to_string()
}

/// Like build, but you can specify the location where files should be generated.
/// Not generally recommended for use in build scripts.
pub(crate) fn build_to_custom_directory<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    custom_gendir: Option<PathBuf>,
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderResult
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    rust_version_check();
    let gen_location_strategy = match custom_gendir {
        None => FileLocationStrategy::new(),
        Some(custom_dir) => FileLocationStrategy::Custom(custom_dir),
    };
    let incdir = gen_location_strategy.get_include_dir();
    ensure_created(&incdir)?;
    let cxxdir = gen_location_strategy.get_cxx_dir();
    ensure_created(&cxxdir)?;
    let rsdir = gen_location_strategy.get_rs_dir();
    ensure_created(&rsdir)?;
    // We are incredibly unsophisticated in our directory arrangement here
    // compared to cxx. I have no doubt that we will need to replicate just
    // about everything cxx does, in due course...
    // Write cxx.h to that location, as it may be needed by
    // some of our generated code.
    write_to_file(&incdir, "cxx.h", crate::HEADER.as_bytes())?;

    let autocxx_inc = build_autocxx_inc(autocxx_incs, &incdir);
    // pass on th
    gen_location_strategy.set_cargo_env_vars_for_build();

    let mut parsed_file = crate::parse_file(rs_file).map_err(BuilderError::ParseError)?;
    parsed_file
        .resolve_all(autocxx_inc, extra_clang_args, dependency_recorder)
        .map_err(BuilderError::ParseError)?;
    build_with_existing_parsed_file(parsed_file, cxxdir, incdir, rsdir)
}

pub(crate) fn build_with_existing_parsed_file(
    parsed_file: ParsedFile,
    cxxdir: PathBuf,
    incdir: PathBuf,
    rsdir: PathBuf,
) -> BuilderResult {
    let mut counter = 0;
    let mut builder = cc::Build::new();
    builder.cpp(true);
    let mut generated_rs = Vec::new();
    for include_cpp in parsed_file.get_autocxxes() {
        for inc_dir in include_cpp.include_dirs() {
            builder.include(inc_dir);
        }
        let generated_code = include_cpp
            .generate_h_and_cxx()
            .map_err(BuilderError::InvalidCxx)?;
        for filepair in generated_code.0 {
            let fname = format!("gen{}.cxx", counter);
            counter += 1;
            let gen_cxx_path = write_to_file(&cxxdir, &fname, &filepair.implementation)?;
            builder.file(gen_cxx_path);

            write_to_file(&incdir, &filepair.header_name, &filepair.header)?;
            let fname = include_cpp.get_rs_filename();
            let rs = include_cpp.generate_rs();
            generated_rs.push(write_rs_to_file(&rsdir, &fname, rs)?);
        }
    }
    if counter == 0 {
        Err(BuilderError::NoIncludeCxxMacrosFound)
    } else {
        Ok(BuilderSuccess(builder, generated_rs))
    }
}

fn ensure_created(dir: &Path) -> Result<(), BuilderError> {
    std::fs::create_dir_all(dir)
        .map_err(|e| BuilderError::UnableToCreateDirectory(e, dir.to_path_buf()))
}

fn build_autocxx_inc<I, T>(paths: I, extra_path: &Path) -> Vec<PathBuf>
where
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    paths
        .into_iter()
        .map(|p| PathBuf::from(p.as_ref()))
        .chain(std::iter::once(extra_path.to_path_buf()))
        .collect()
}

fn write_to_file(dir: &Path, filename: &str, content: &[u8]) -> Result<PathBuf, BuilderError> {
    let path = dir.join(filename);
    try_write_to_file(&path, content).map_err(|e| BuilderError::FileWriteFail(e, path.clone()))?;
    Ok(path)
}

fn try_write_to_file(path: &Path, content: &[u8]) -> std::io::Result<()> {
    let mut f = File::create(path)?;
    f.write_all(content)
}

fn write_rs_to_file(
    dir: &Path,
    filename: &str,
    content: TokenStream,
) -> Result<PathBuf, BuilderError> {
    write_to_file(dir, filename, content.to_string().as_bytes())
}

fn rust_version_check() {
    if !version_check::is_min_version("1.48.0").unwrap_or(false) {
        panic!("Rust 1.48 or later is required.")
    }
}
<commit_msg>Small change.<commit_after>// Copyright 2020 Google LLC
//
// 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
//
//    https://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 autocxx_parser::file_locations::FileLocationStrategy;
use proc_macro2::TokenStream;

use crate::{ParseError, ParsedFile, RebuildDependencyRecorder};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::{ffi::OsStr, io, process};
use std::{fmt::Display, fs::File};

/// Errors returned during creation of a cc::Build from an include_cxx
/// macro.
#[derive(Debug)]
pub enum BuilderError {
    /// The cxx module couldn't parse the code generated by autocxx.
    /// This could well be a bug in autocxx.
    InvalidCxx(cxx_gen::Error),
    /// The .rs file didn't exist or couldn't be parsed.
    ParseError(ParseError),
    /// We couldn't write the c++ code to disk.
    FileWriteFail(std::io::Error, PathBuf),
    /// No `include_cxx` macro was found anywhere.
    NoIncludeCxxMacrosFound,
    /// Unable to create one of the directories to which we need to write
    UnableToCreateDirectory(std::io::Error, PathBuf),
}

impl Display for BuilderError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BuilderError::ParseError(pe) => write!(f, "Unable to parse .rs file: {}", pe)?,
            BuilderError::InvalidCxx(ee) => write!(f, "cxx was unable to understand the code generated by autocxx (likely a bug in autocxx; please report.) {}", ee)?,
            BuilderError::FileWriteFail(ee, pb) => write!(f, "Unable to write to {}: {}", pb.to_string_lossy(), ee)?,
            BuilderError::NoIncludeCxxMacrosFound => write!(f, "No include_cpp! macro found")?,
            BuilderError::UnableToCreateDirectory(ee, pb) => write!(f, "Unable to create directory {}: {}", pb.to_string_lossy(), ee)?,
        }
        Ok(())
    }
}

pub type BuilderBuild = cc::Build;

pub struct BuilderSuccess(pub BuilderBuild, pub Vec<PathBuf>);

/// Results of a build.
pub type BuilderResult = Result<BuilderSuccess, BuilderError>;

/// Build autocxx C++ files and return a cc::Build you can use to build
/// more from a build.rs file.
/// You need to provide the Rust file path and the iterator of paths
/// which should be used as include directories.
pub fn build<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderResult
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    build_to_custom_directory(
        rs_file,
        autocxx_incs,
        extra_clang_args,
        None,
        dependency_recorder,
    )
}

/// Builds successfully, or exits the process displaying a suitable
/// message.
pub fn expect_build<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderSuccess
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    build(rs_file, autocxx_incs, extra_clang_args, dependency_recorder).unwrap_or_else(|err| {
        let _ = writeln!(io::stderr(), "\n\nautocxx error: {}\n\n", err);
        process::exit(1);
    })
}

/// Like build, but you can specify the location where files should be generated.
/// Not generally recommended for use in build scripts.
pub(crate) fn build_to_custom_directory<P1, I, T>(
    rs_file: P1,
    autocxx_incs: I,
    extra_clang_args: &[&str],
    custom_gendir: Option<PathBuf>,
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
) -> BuilderResult
where
    P1: AsRef<Path>,
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    rust_version_check();
    let gen_location_strategy = match custom_gendir {
        None => FileLocationStrategy::new(),
        Some(custom_dir) => FileLocationStrategy::Custom(custom_dir),
    };
    let incdir = gen_location_strategy.get_include_dir();
    ensure_created(&incdir)?;
    let cxxdir = gen_location_strategy.get_cxx_dir();
    ensure_created(&cxxdir)?;
    let rsdir = gen_location_strategy.get_rs_dir();
    ensure_created(&rsdir)?;
    // We are incredibly unsophisticated in our directory arrangement here
    // compared to cxx. I have no doubt that we will need to replicate just
    // about everything cxx does, in due course...
    // Write cxx.h to that location, as it may be needed by
    // some of our generated code.
    write_to_file(&incdir, "cxx.h", crate::HEADER.as_bytes())?;

    let autocxx_inc = build_autocxx_inc(autocxx_incs, &incdir);
    // pass on th
    gen_location_strategy.set_cargo_env_vars_for_build();

    let mut parsed_file = crate::parse_file(rs_file).map_err(BuilderError::ParseError)?;
    parsed_file
        .resolve_all(autocxx_inc, extra_clang_args, dependency_recorder)
        .map_err(BuilderError::ParseError)?;
    build_with_existing_parsed_file(parsed_file, cxxdir, incdir, rsdir)
}

pub(crate) fn build_with_existing_parsed_file(
    parsed_file: ParsedFile,
    cxxdir: PathBuf,
    incdir: PathBuf,
    rsdir: PathBuf,
) -> BuilderResult {
    let mut counter = 0;
    let mut builder = cc::Build::new();
    builder.cpp(true);
    let mut generated_rs = Vec::new();
    for include_cpp in parsed_file.get_autocxxes() {
        for inc_dir in include_cpp.include_dirs() {
            builder.include(inc_dir);
        }
        let generated_code = include_cpp
            .generate_h_and_cxx()
            .map_err(BuilderError::InvalidCxx)?;
        for filepair in generated_code.0 {
            let fname = format!("gen{}.cxx", counter);
            counter += 1;
            let gen_cxx_path = write_to_file(&cxxdir, &fname, &filepair.implementation)?;
            builder.file(gen_cxx_path);

            write_to_file(&incdir, &filepair.header_name, &filepair.header)?;
            let fname = include_cpp.get_rs_filename();
            let rs = include_cpp.generate_rs();
            generated_rs.push(write_rs_to_file(&rsdir, &fname, rs)?);
        }
    }
    if counter == 0 {
        Err(BuilderError::NoIncludeCxxMacrosFound)
    } else {
        Ok(BuilderSuccess(builder, generated_rs))
    }
}

fn ensure_created(dir: &Path) -> Result<(), BuilderError> {
    std::fs::create_dir_all(dir)
        .map_err(|e| BuilderError::UnableToCreateDirectory(e, dir.to_path_buf()))
}

fn build_autocxx_inc<I, T>(paths: I, extra_path: &Path) -> Vec<PathBuf>
where
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    paths
        .into_iter()
        .map(|p| PathBuf::from(p.as_ref()))
        .chain(std::iter::once(extra_path.to_path_buf()))
        .collect()
}

fn write_to_file(dir: &Path, filename: &str, content: &[u8]) -> Result<PathBuf, BuilderError> {
    let path = dir.join(filename);
    try_write_to_file(&path, content).map_err(|e| BuilderError::FileWriteFail(e, path.clone()))?;
    Ok(path)
}

fn try_write_to_file(path: &Path, content: &[u8]) -> std::io::Result<()> {
    let mut f = File::create(path)?;
    f.write_all(content)
}

fn write_rs_to_file(
    dir: &Path,
    filename: &str,
    content: TokenStream,
) -> Result<PathBuf, BuilderError> {
    write_to_file(dir, filename, content.to_string().as_bytes())
}

fn rust_version_check() {
    if !version_check::is_min_version("1.48.0").unwrap_or(false) {
        panic!("Rust 1.48 or later is required.")
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::str::FromStr;

use yaml::Yaml;

pub struct Parameter {
    pub description: Option<String>,
    pub display_name: Option<String>,
    pub name: String,
    pub parameter_type: Option<ParameterType>,
    pub required: bool,
    pub value: Option<String>,
}

pub enum ParameterType {
    Bool,
    Int,
    String,
}

pub type ParamMap = HashMap<String, Parameter>;
pub type UserValues = HashMap<String, String>;

impl Parameter {
    pub fn new(yaml: &Yaml, user_values: &UserValues) -> Result<Self, String> {
        let description = match yaml["description"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let display_name = match yaml["displayName"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let name = match yaml["name"] {
            Yaml::String(ref name) => name.clone(),
            _ => return Err("Parameters must have a \"name\" field.".to_owned()),
        };
        let parameter_type = match yaml["parameterType"].as_str() {
            Some(ref parameter_type) => Some(try!(parameter_type.parse())),
            None => None,
        };
        let required = yaml["required"].as_bool().unwrap_or(false);
        let value = match user_values.get(&name) {
            Some(value) => Some(value.clone()),
            None => match yaml["value"] {
                Yaml::Boolean(ref value)  => Some(format!("{}", value)),
                Yaml::Integer(ref value) => Some(format!("{}", value)),
                Yaml::String(ref value) => Some(value.clone()),
                _ => if required {
                    return Err(
                        format!(
                            "Parameter {} required and must be boolean, i64, or string.",
                            display_name.unwrap_or(name),
                        )
                    )
                } else {
                    None
                },
            }
        };

        Ok(Parameter {
            description: description,
            display_name: display_name,
            name: name,
            parameter_type: parameter_type,
            required: required,
            value: value,
        })
    }
}

impl FromStr for ParameterType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "base64" => Ok(ParameterType::String),
            "bool" => Ok(ParameterType::Bool),
            "int" => Ok(ParameterType::Int),
            "string" => Ok(ParameterType::String),
            _ => Err("parameterType must be base64, bool, int, or string.".to_owned()),
        }
    }
}

pub fn user_values(parameters: Vec<String>) -> Result<UserValues, String> {
    let mut user_values = HashMap::new();

    for parameter in parameters {
        let mut parts: Vec<String> = parameter.split('=').map(|s| s.to_string()).collect();

        if parts.len() == 2 {
            user_values.insert(parts.remove(0), parts.remove(0));
        } else {
            return Err("Parameters must be supplied in the form KEY=VALUE.".to_string());
        }
    }

    Ok(user_values)
}
<commit_msg>if parameter type is known, let it be known<commit_after>use std::collections::HashMap;
use std::str::FromStr;

use yaml::Yaml;

pub struct Parameter {
    pub description: Option<String>,
    pub display_name: Option<String>,
    pub name: String,
    pub parameter_type: Option<ParameterType>,
    pub required: bool,
    pub value: Option<String>,
}

pub enum ParameterType {
    Bool,
    Int,
    String,
}

pub type ParamMap = HashMap<String, Parameter>;
pub type UserValues = HashMap<String, String>;

impl Parameter {
    pub fn new(yaml: &Yaml, user_values: &UserValues) -> Result<Self, String> {
        let description = match yaml["description"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let display_name = match yaml["displayName"] {
            Yaml::String(ref description) => Some(description.clone()),
            _ => None,
        };
        let name = match yaml["name"] {
            Yaml::String(ref name) => name.clone(),
            _ => return Err("Parameters must have a \"name\" field.".to_owned()),
        };
        let parameter_type = match yaml["parameterType"].as_str() {
            Some(ref parameter_type) => Some(try!(parameter_type.parse())),
            None => None,
        };
        let required = yaml["required"].as_bool().unwrap_or(false);
        let value = match user_values.get(&name) {
            Some(value) => Some(value.clone()),
            None => match yaml["value"] {
                Yaml::Boolean(ref value)  => Some(format!("{}", value)),
                Yaml::Integer(ref value) => Some(format!("{}", value)),
                Yaml::String(ref value) => Some(value.clone()),
                _ => if required {
                    return Err(
                        format!(
                            "Parameter {} required and must be {}",
                            display_name.unwrap_or(name),
                            parameter_type.map(|pt| match pt {
                                ParameterType::Bool => "a bool",
                                ParameterType::Int => "an int",
                                ParameterType::String => "a string"
                            }).unwrap_or("base64, bool, int, or string")
                        )
                    )
                } else {
                    None
                },
            }
        };

        Ok(Parameter {
            description: description,
            display_name: display_name,
            name: name,
            parameter_type: parameter_type,
            required: required,
            value: value,
        })
    }
}

impl FromStr for ParameterType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "base64" => Ok(ParameterType::String),
            "bool" => Ok(ParameterType::Bool),
            "int" => Ok(ParameterType::Int),
            "string" => Ok(ParameterType::String),
            _ => Err("parameterType must be base64, bool, int, or string.".to_owned()),
        }
    }
}

pub fn user_values(parameters: Vec<String>) -> Result<UserValues, String> {
    let mut user_values = HashMap::new();

    for parameter in parameters {
        let mut parts: Vec<String> = parameter.split('=').map(|s| s.to_string()).collect();

        if parts.len() == 2 {
            user_values.insert(parts.remove(0), parts.remove(0));
        } else {
            return Err("Parameters must be supplied in the form KEY=VALUE.".to_string());
        }
    }

    Ok(user_values)
}
<|endoftext|>


<commit_before>use scroll;
use pe::error;
use pe::section_table;
use pe::utils;
use pe::data_directories;

#[derive(Debug, Clone)]
pub struct HintNameTableEntry {
    pub hint: u16,
    pub name: String,
}

impl HintNameTableEntry {
    fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, mut offset: usize) -> error::Result<Self> {
        let mut offset = &mut offset;
        let hint = bytes.gread(offset, scroll::LE)?;
        let name = bytes.pread::<&str>(*offset, 0)?.to_string();
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry {
    OrdinalNumber(u16),
    HintNameTableRVA ((u32, HintNameTableEntry)), // [u8; 31] bitfield :/
}

#[derive(Debug)]
pub struct ImportLookupTableEntry {
    pub bitfield: u32,
    pub synthetic: SyntheticImportLookupTableEntry,
}

pub type ImportLookupTable = Vec<ImportLookupTableEntry>;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;

impl ImportLookupTableEntry {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable> {
        let le = scroll::LE;
        let mut offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: u32 = bytes.gread(offset, le)?;
            if bitfield == 0 {
                println!("imports done");
                break;
            } else {
                let synthetic = {
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 {
                        let ordinal = (0xffff & bitfield) as u16;
                        println!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield & IMPORT_RVA_MASK_32;
                        let hentry = {
                            println!("searching for RVA {:#x}", rva);
                            let offset = utils::find_offset(rva as usize, sections).unwrap();
                            println!("offset {:#x}", offset);
                            HintNameTableEntry::parse(bytes, offset)?
                            //HintNameTableEntry {hint = 0; name = "".to_string()}
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                let entry = ImportLookupTableEntry { bitfield: bitfield, synthetic: synthetic };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<u32>;

pub const SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY: usize = 4;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: String,
    /// Computed
    pub import_lookup_table: ImportLookupTable,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl SyntheticImportDirectoryEntry {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<Self> {
        let le = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?.to_string();
        let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
        let import_lookup_table_offset = utils::find_offset(import_lookup_table_rva as usize, sections).unwrap();
        let import_lookup_table = ImportLookupTableEntry::parse(bytes, import_lookup_table_offset, sections)?;
        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).unwrap();
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread(import_address_table_offset, le)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
pub struct ImportData {
    pub import_data: Vec<SyntheticImportDirectoryEntry>,
}

impl ImportData {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<Self> {
        let import_directory_table_rva = dd.virtual_address as usize;
        let mut offset = &mut utils::find_offset(import_directory_table_rva, sections).unwrap();
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread(offset, scroll::LE)?;
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse(bytes, import_directory_entry, sections)?;
                import_data.push(entry);
            }
        }
        println!("finished import directory table");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
pub struct Import {
    pub name: String,
    pub dll: String,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl Import {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(_bytes: &B, import_data: &ImportData, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Self>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            let import_lookup_table = &data.import_lookup_table;
            let dll = data.name.to_owned();
            let import_base = data.import_directory_entry.import_address_table_rva as usize;
            println!("getting imports from {}", &dll);
            for (i, entry) in import_lookup_table.iter().enumerate() {
                let offset = import_base + (i * SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY);
                use self::SyntheticImportLookupTableEntry::*;
                let (rva, name, ordinal) =
                    match &entry.synthetic {
                        &HintNameTableRVA ((rva, ref hint_entry)) => {
                            let res = (rva, hint_entry.name.to_owned(), hint_entry.hint.clone());
                            // if hint_entry.name = "" && hint_entry.hint = 0 {
                            //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                            // }
                            res
                        },
                        &OrdinalNumber(ordinal) => {
                            let name = format!("ORDINAL {}", ordinal);
                            (0x0, name, ordinal)
                        }
                    };
                let import =
                    Import {
                        name: name,
                        ordinal: ordinal, dll: dll.to_owned(),
                        size: 4, offset: offset, rva: rva as usize
                    };
                imports.push(import);
            }
        }
        Ok (imports)
    }
}
<commit_msg>pe.import: comment printlns<commit_after>use scroll;
use pe::error;
use pe::section_table;
use pe::utils;
use pe::data_directories;

#[derive(Debug, Clone)]
pub struct HintNameTableEntry {
    pub hint: u16,
    pub name: String,
}

impl HintNameTableEntry {
    fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, mut offset: usize) -> error::Result<Self> {
        let mut offset = &mut offset;
        let hint = bytes.gread(offset, scroll::LE)?;
        let name = bytes.pread::<&str>(*offset, 0)?.to_string();
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry {
    OrdinalNumber(u16),
    HintNameTableRVA ((u32, HintNameTableEntry)), // [u8; 31] bitfield :/
}

#[derive(Debug)]
pub struct ImportLookupTableEntry {
    pub bitfield: u32,
    pub synthetic: SyntheticImportLookupTableEntry,
}

pub type ImportLookupTable = Vec<ImportLookupTableEntry>;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;

impl ImportLookupTableEntry {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable> {
        let le = scroll::LE;
        let mut offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: u32 = bytes.gread(offset, le)?;
            if bitfield == 0 {
                //println!("imports done");
                break;
            } else {
                let synthetic = {
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 {
                        let ordinal = (0xffff & bitfield) as u16;
                        //println!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield & IMPORT_RVA_MASK_32;
                        let hentry = {
                            //println!("searching for RVA {:#x}", rva);
                            let offset = utils::find_offset(rva as usize, sections).unwrap();
                            //println!("offset {:#x}", offset);
                            HintNameTableEntry::parse(bytes, offset)?
                            //HintNameTableEntry {hint = 0; name = "".to_string()}
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                let entry = ImportLookupTableEntry { bitfield: bitfield, synthetic: synthetic };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<u32>;

pub const SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY: usize = 4;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: String,
    /// Computed
    pub import_lookup_table: ImportLookupTable,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl SyntheticImportDirectoryEntry {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<Self> {
        let le = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?.to_string();
        let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
        let import_lookup_table_offset = utils::find_offset(import_lookup_table_rva as usize, sections).unwrap();
        let import_lookup_table = ImportLookupTableEntry::parse(bytes, import_lookup_table_offset, sections)?;
        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).unwrap();
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread(import_address_table_offset, le)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
pub struct ImportData {
    pub import_data: Vec<SyntheticImportDirectoryEntry>,
}

impl ImportData {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(bytes: &B, dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<Self> {
        let import_directory_table_rva = dd.virtual_address as usize;
        let mut offset = &mut utils::find_offset(import_directory_table_rva, sections).unwrap();
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread(offset, scroll::LE)?;
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse(bytes, import_directory_entry, sections)?;
                import_data.push(entry);
            }
        }
        //println!("finished import directory table");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
pub struct Import {
    pub name: String,
    pub dll: String,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl Import {
    pub fn parse<B: scroll::Gread + scroll::Gread<scroll::Error, u8>>(_bytes: &B, import_data: &ImportData, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Self>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            let import_lookup_table = &data.import_lookup_table;
            let dll = data.name.to_owned();
            let import_base = data.import_directory_entry.import_address_table_rva as usize;
            //println!("getting imports from {}", &dll);
            for (i, entry) in import_lookup_table.iter().enumerate() {
                let offset = import_base + (i * SIZEOF_IMPORT_ADDRESS_TABLE_ENTRY);
                use self::SyntheticImportLookupTableEntry::*;
                let (rva, name, ordinal) =
                    match &entry.synthetic {
                        &HintNameTableRVA ((rva, ref hint_entry)) => {
                            let res = (rva, hint_entry.name.to_owned(), hint_entry.hint.clone());
                            // if hint_entry.name = "" && hint_entry.hint = 0 {
                            //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                            // }
                            res
                        },
                        &OrdinalNumber(ordinal) => {
                            let name = format!("ORDINAL {}", ordinal);
                            (0x0, name, ordinal)
                        }
                    };
                let import =
                    Import {
                        name: name,
                        ordinal: ordinal, dll: dll.to_owned(),
                        size: 4, offset: offset, rva: rva as usize
                    };
                imports.push(import);
            }
        }
        Ok (imports)
    }
}
<|endoftext|>


<commit_before>use std::fs;
use std::io;
use std::path::Path;

#[cfg(unix)]
use tokio::net::UnixListener;
#[cfg(unix)]
use tokio::net::UnixStream;

use futures::stream;

use futures::stream::Stream;
use futures::Future;

use crate::assert_types::assert_send_stream;
use crate::net::addr::AnySocketAddr;
use crate::net::connect::ToClientStream;
use crate::net::listen::ToServerStream;
use crate::net::listen::ToSocketListener;
use crate::net::listen::ToTokioListener;
use crate::net::socket::SocketStream;
use crate::ServerConf;
use std::fmt;
#[cfg(unix)]
use std::os::unix::net::SocketAddr;
use std::path::PathBuf;
use std::pin::Pin;
use tokio::runtime::Handle;

/// Unix socket address, which is filesystem path.
///
/// Note although this type is available on Windows, unix sockets don't work on Windows.
#[derive(Debug, Clone, PartialEq)]
pub struct SocketAddrUnix(pub(crate) PathBuf);

impl From<PathBuf> for SocketAddrUnix {
    fn from(p: PathBuf) -> Self {
        SocketAddrUnix(p)
    }
}

impl From<&Path> for SocketAddrUnix {
    fn from(p: &Path) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl From<&str> for SocketAddrUnix {
    fn from(p: &str) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl From<String> for SocketAddrUnix {
    fn from(p: String) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl fmt::Display for SocketAddrUnix {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0.display(), f)
    }
}

#[cfg(unix)]
impl From<SocketAddr> for SocketAddrUnix {
    fn from(s: SocketAddr) -> Self {
        // can be unnamed
        SocketAddrUnix(s.as_pathname().unwrap_or(Path::new("")).to_owned())
    }
}

impl ToSocketListener for SocketAddrUnix {
    #[cfg(unix)]
    fn to_listener(&self, _conf: &ServerConf) -> io::Result<Box<dyn ToTokioListener + Send>> {
        debug!("binding socket to {}", self);
        Ok(Box::new(::std::os::unix::net::UnixListener::bind(&self.0)?))
    }

    #[cfg(not(unix))]
    fn to_listener(&self, _conf: &ServerConf) -> io::Result<Box<ToTokioListener + Send>> {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "cannot use unix sockets on non-unix",
        ))
    }

    fn cleanup(&self) {
        if self.0.exists() {
            fs::remove_file(&self.0).expect("removing socket during shutdown");
        }
    }
}

#[cfg(unix)]
impl ToTokioListener for ::std::os::unix::net::UnixListener {
    fn to_tokio_listener(self: Box<Self>, handle: &Handle) -> Box<dyn ToServerStream> {
        handle.enter(|| Box::new(UnixListener::from_std(*self).unwrap()))
    }

    fn local_addr(&self) -> io::Result<AnySocketAddr> {
        let addr = self.local_addr().unwrap();
        let path = addr.as_pathname().unwrap();
        Ok(AnySocketAddr::Unix(SocketAddrUnix::from(path)))
    }
}

#[cfg(unix)]
impl ToServerStream for UnixListener {
    fn incoming(
        self: Box<Self>,
    ) -> Pin<
        Box<
            dyn Stream<Item = io::Result<(Pin<Box<dyn SocketStream + Send>>, AnySocketAddr)>>
                + Send,
        >,
    > {
        let unix_listener = *self;

        let stream = stream::unfold(unix_listener, |mut unix_listener_listener| async {
            let r = match unix_listener_listener.accept().await {
                Ok((socket, addr)) => Ok((
                    Box::pin(socket) as Pin<Box<dyn SocketStream + Send>>,
                    AnySocketAddr::Unix(addr.into()),
                )),
                Err(e) => Err(e),
            };
            Some((r, unix_listener_listener))
        });

        let stream = assert_send_stream::<
            io::Result<(Pin<Box<dyn SocketStream + Send>>, AnySocketAddr)>,
            _,
        >(stream);

        Box::pin(stream)
    }
}

impl ToClientStream for SocketAddrUnix {
    #[cfg(unix)]
    fn connect(
        &self,
        handle: &Handle,
    ) -> Pin<Box<dyn Future<Output = io::Result<Pin<Box<dyn SocketStream + Send>>>> + Send>> {
        // TODO: async connect
        let stream = match std::os::unix::net::UnixStream::connect(&self.0) {
            Ok(stream) => stream,
            Err(e) => return Box::pin(async { Err(e) }),
        };
        match handle.enter(|| UnixStream::from_std(stream)) {
            Ok(stream) => {
                Box::pin(async { Ok(Box::pin(stream) as Pin<Box<dyn SocketStream + Send>>) })
            }
            Err(e) => return Box::pin(async { Err(e) }),
        }
    }

    #[cfg(not(unix))]
    fn connect(
        &self,
        _handle: &Handle,
    ) -> Pin<Box<dyn Future<Output = io::Result<Pin<Box<dyn SocketStream + Send>>>> + Send>> {
        use futures::future;
        Box::pin(future::err(io::Error::new(
            io::ErrorKind::Other,
            "cannot use unix sockets on non-unix",
        )))
    }

    fn socket_addr(&self) -> AnySocketAddr {
        AnySocketAddr::Unix(self.clone())
    }
}

#[cfg(unix)]
impl SocketStream for UnixStream {
    fn is_tcp(&self) -> bool {
        false
    }

    fn set_nodelay(&self, _no_delay: bool) -> io::Result<()> {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "Cannot set nodelay on unix domain socket",
        ))
    }
}
<commit_msg>Missing dyn warning<commit_after>use std::fs;
use std::io;
use std::path::Path;

#[cfg(unix)]
use tokio::net::UnixListener;
#[cfg(unix)]
use tokio::net::UnixStream;

use futures::stream;

use futures::stream::Stream;
use futures::Future;

use crate::assert_types::assert_send_stream;
use crate::net::addr::AnySocketAddr;
use crate::net::connect::ToClientStream;
use crate::net::listen::ToServerStream;
use crate::net::listen::ToSocketListener;
use crate::net::listen::ToTokioListener;
use crate::net::socket::SocketStream;
use crate::ServerConf;
use std::fmt;
#[cfg(unix)]
use std::os::unix::net::SocketAddr;
use std::path::PathBuf;
use std::pin::Pin;
use tokio::runtime::Handle;

/// Unix socket address, which is filesystem path.
///
/// Note although this type is available on Windows, unix sockets don't work on Windows.
#[derive(Debug, Clone, PartialEq)]
pub struct SocketAddrUnix(pub(crate) PathBuf);

impl From<PathBuf> for SocketAddrUnix {
    fn from(p: PathBuf) -> Self {
        SocketAddrUnix(p)
    }
}

impl From<&Path> for SocketAddrUnix {
    fn from(p: &Path) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl From<&str> for SocketAddrUnix {
    fn from(p: &str) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl From<String> for SocketAddrUnix {
    fn from(p: String) -> Self {
        SocketAddrUnix(p.into())
    }
}

impl fmt::Display for SocketAddrUnix {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0.display(), f)
    }
}

#[cfg(unix)]
impl From<SocketAddr> for SocketAddrUnix {
    fn from(s: SocketAddr) -> Self {
        // can be unnamed
        SocketAddrUnix(s.as_pathname().unwrap_or(Path::new("")).to_owned())
    }
}

impl ToSocketListener for SocketAddrUnix {
    #[cfg(unix)]
    fn to_listener(&self, _conf: &ServerConf) -> io::Result<Box<dyn ToTokioListener + Send>> {
        debug!("binding socket to {}", self);
        Ok(Box::new(::std::os::unix::net::UnixListener::bind(&self.0)?))
    }

    #[cfg(not(unix))]
    fn to_listener(&self, _conf: &ServerConf) -> io::Result<Box<dyn ToTokioListener + Send>> {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "cannot use unix sockets on non-unix",
        ))
    }

    fn cleanup(&self) {
        if self.0.exists() {
            fs::remove_file(&self.0).expect("removing socket during shutdown");
        }
    }
}

#[cfg(unix)]
impl ToTokioListener for ::std::os::unix::net::UnixListener {
    fn to_tokio_listener(self: Box<Self>, handle: &Handle) -> Box<dyn ToServerStream> {
        handle.enter(|| Box::new(UnixListener::from_std(*self).unwrap()))
    }

    fn local_addr(&self) -> io::Result<AnySocketAddr> {
        let addr = self.local_addr().unwrap();
        let path = addr.as_pathname().unwrap();
        Ok(AnySocketAddr::Unix(SocketAddrUnix::from(path)))
    }
}

#[cfg(unix)]
impl ToServerStream for UnixListener {
    fn incoming(
        self: Box<Self>,
    ) -> Pin<
        Box<
            dyn Stream<Item = io::Result<(Pin<Box<dyn SocketStream + Send>>, AnySocketAddr)>>
                + Send,
        >,
    > {
        let unix_listener = *self;

        let stream = stream::unfold(unix_listener, |mut unix_listener_listener| async {
            let r = match unix_listener_listener.accept().await {
                Ok((socket, addr)) => Ok((
                    Box::pin(socket) as Pin<Box<dyn SocketStream + Send>>,
                    AnySocketAddr::Unix(addr.into()),
                )),
                Err(e) => Err(e),
            };
            Some((r, unix_listener_listener))
        });

        let stream = assert_send_stream::<
            io::Result<(Pin<Box<dyn SocketStream + Send>>, AnySocketAddr)>,
            _,
        >(stream);

        Box::pin(stream)
    }
}

impl ToClientStream for SocketAddrUnix {
    #[cfg(unix)]
    fn connect(
        &self,
        handle: &Handle,
    ) -> Pin<Box<dyn Future<Output = io::Result<Pin<Box<dyn SocketStream + Send>>>> + Send>> {
        // TODO: async connect
        let stream = match std::os::unix::net::UnixStream::connect(&self.0) {
            Ok(stream) => stream,
            Err(e) => return Box::pin(async { Err(e) }),
        };
        match handle.enter(|| UnixStream::from_std(stream)) {
            Ok(stream) => {
                Box::pin(async { Ok(Box::pin(stream) as Pin<Box<dyn SocketStream + Send>>) })
            }
            Err(e) => return Box::pin(async { Err(e) }),
        }
    }

    #[cfg(not(unix))]
    fn connect(
        &self,
        _handle: &Handle,
    ) -> Pin<Box<dyn Future<Output = io::Result<Pin<Box<dyn SocketStream + Send>>>> + Send>> {
        use futures::future;
        Box::pin(future::err(io::Error::new(
            io::ErrorKind::Other,
            "cannot use unix sockets on non-unix",
        )))
    }

    fn socket_addr(&self) -> AnySocketAddr {
        AnySocketAddr::Unix(self.clone())
    }
}

#[cfg(unix)]
impl SocketStream for UnixStream {
    fn is_tcp(&self) -> bool {
        false
    }

    fn set_nodelay(&self, _no_delay: bool) -> io::Result<()> {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "Cannot set nodelay on unix domain socket",
        ))
    }
}
<|endoftext|>