<commit_before>// Copyright (c) 2014-2016 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 sending and receiving data link layer packets using the WinPcap library

extern crate libc;
extern crate winapi;

use bindings::{bpf, winpcap};
use {DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface};

use ipnetwork::{ip_mask_to_prefix, IpNetwork};

use std::cmp;
use std::collections::VecDeque;
use std::ffi::{CStr, CString};
use std::io;
use std::mem;
use std::slice;
use std::str::from_utf8_unchecked;
use std::sync::Arc;

struct WinPcapAdapter {
    adapter: winpcap::LPADAPTER,
}

impl Drop for WinPcapAdapter {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketCloseAdapter(self.adapter);
        }
    }
}

struct WinPcapPacket {
    packet: winpcap::LPPACKET,
}

impl Drop for WinPcapPacket {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketFreePacket(self.packet);
        }
    }
}

/// WinPcap specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
        }
    }
}

/// Create a datalink channel using the WinPcap library
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<super::Channel> {
    let mut read_buffer = Vec::new();
    read_buffer.resize(config.read_buffer_size, 0u8);

    let mut write_buffer = Vec::new();
    write_buffer.resize(config.write_buffer_size, 0u8);

    let adapter = unsafe {
        let net_if_str = CString::new(network_interface.name.as_bytes()).unwrap();
        winpcap::PacketOpenAdapter(net_if_str.as_ptr() as *mut libc::c_char)
    };
    if adapter.is_null() {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetHwFilter(adapter, winpcap::NDIS_PACKET_TYPE_PROMISCUOUS) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Set kernel buffer size
    let ret = unsafe { winpcap::PacketSetBuff(adapter, config.read_buffer_size as libc::c_int) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Immediate mode
    let ret = unsafe { winpcap::PacketSetMinToCopy(adapter, 1) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let read_packet = unsafe { winpcap::PacketAllocatePacket() };
    if read_packet.is_null() {
        unsafe {
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(read_packet,
                                  read_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.read_buffer_size as winpcap::UINT)
    }

    let write_packet = unsafe { winpcap::PacketAllocatePacket() };
    if write_packet.is_null() {
        unsafe {
            winpcap::PacketFreePacket(read_packet);
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(write_packet,
                                  write_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.write_buffer_size as winpcap::UINT)
    }

    let adapter = Arc::new(WinPcapAdapter { adapter: adapter });
    let sender = Box::new(DataLinkSenderImpl {
        adapter: adapter.clone(),
        _write_buffer: write_buffer,
        packet: WinPcapPacket { packet: write_packet },
    });
    let receiver = Box::new(DataLinkReceiverImpl {
        adapter: adapter,
        _read_buffer: read_buffer,
        packet: WinPcapPacket { packet: read_packet },
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(unsafe { (*read_packet).Length } as usize / 64),
    });
    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    adapter: Arc<WinPcapAdapter>,
    _write_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= unsafe { (*self.packet.packet).Length } as usize {
            None
        } else {
            let min = unsafe { cmp::min((*self.packet.packet).Length as usize, len) };
            let slice: &mut [u8] =
                unsafe { slice::from_raw_parts_mut((*self.packet.packet).Buffer as *mut u8, min) };
            for chunk in slice.chunks_mut(packet_size) {
                func(chunk);

                // Make sure the right length of packet is sent
                let old_len = unsafe { (*self.packet.packet).Length };
                unsafe {
                    (*self.packet.packet).Length = packet_size as u32;
                }

                let ret = unsafe {
                    winpcap::PacketSendPacket(self.adapter.adapter, self.packet.packet, 0)
                };

                unsafe {
                    (*self.packet.packet).Length = old_len;
                }

                if ret == 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        self.build_and_send(1,
                            packet.len(),
                            &mut |eh: &mut [u8]| {
                                eh.copy_from_slice(packet);
                            })
    }
}

unsafe impl Send for DataLinkSenderImpl {}
unsafe impl Sync for DataLinkSenderImpl {}

struct DataLinkReceiverImpl {
    adapter: Arc<WinPcapAdapter>,
    _read_buffer: Vec<u8>,
    packet: WinPcapPacket,
    packets: VecDeque<(usize, usize)>,
}

unsafe impl Send for DataLinkReceiverImpl {}
unsafe impl Sync for DataLinkReceiverImpl {}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // NOTE Most of the logic here is identical to FreeBSD/OS X
        while self.packets.is_empty() {
            let ret = unsafe {
                winpcap::PacketReceivePacket(self.adapter.adapter, self.packet.packet, 0)
            };
            let buflen = match ret {
                0 => return Err(io::Error::last_os_error()),
                _ => unsafe { (*self.packet.packet).ulBytesReceived as isize },
            };
            let mut ptr = unsafe { (*self.packet.packet).Buffer };
            let end = unsafe { (*self.packet.packet).Buffer.offset(buflen) };
            while ptr < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                (*self.packet.packet).Buffer as isize;
                    self.packets.push_back((start as usize, (*packet).bh_caplen as usize));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        let slice = unsafe {
            let data = (*self.packet.packet).Buffer as usize + start;
            slice::from_raw_parts(data as *const u8, len)
        };
        Ok(slice)
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use
        // GetAdaptersAddresses
        winpcap::GetAdaptersInfo(&mut tmp, &mut adapters_size);
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    let mut adapters = Vec::with_capacity(vec_size as usize);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as usize);
    while !cursor.is_null() {
        let mac = unsafe {
            MacAddr((*cursor).Address[0],
                    (*cursor).Address[1],
                    (*cursor).Address[2],
                    (*cursor).Address[3],
                    (*cursor).Address[4],
                    (*cursor).Address[5])
        };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING };
        let mut ips = Vec::new();
        while !ip_cursor.is_null() {
            if let Ok(ip_network) = parse_ip_network(ip_cursor) {
                ips.push(ip_network);
            }
            ip_cursor = unsafe { (*ip_cursor).Next };
        }

        unsafe {
            let name_str_ptr = (*cursor).AdapterName.as_ptr() as *const i8;

            let bytes = CStr::from_ptr(name_str_ptr).to_bytes();
            let name_str = from_utf8_unchecked(bytes).to_owned();

            all_ifaces.push(NetworkInterface {
                name: name_str,
                index: (*cursor).Index,
                mac: Some(mac),
                ips: ips,
                // flags: (*cursor).Type, // FIXME [windows]
                flags: 0,
            });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = vec![0u8; 4096];
    let mut buflen = buf.len() as u32;

    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        buf.resize(buflen as usize, 0);

        // Second call should now work with the correct buffer size. If not, this may be
        // due to some privilege or other unforseen issue.
        if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
            panic!("Unable to get interface list despite increasing buffer size");
        }
    }

    let buf_str = unsafe { from_utf8_unchecked(&buf) };
    let iface_names = buf_str.split("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_owned();
                let next = all_ifaces.iter().filter(|x| name[..].ends_with(&x.name[..])).next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        }
        None => (),
    };

    vec
}

fn parse_ip_network(ip_cursor: winpcap::PIP_ADDR_STRING) -> Result<IpNetwork, ()> {
    let ip_str_ptr = unsafe { &(*ip_cursor) }.IpAddress.String.as_ptr() as *const i8;
    let ip_bytes = unsafe { CStr::from_ptr(ip_str_ptr).to_bytes() };
    let ip_str = unsafe { from_utf8_unchecked(ip_bytes).to_owned() };
    let ip = ip_str.parse().map_err(|_| ())?;

    let mask_str_ptr = unsafe { &(*ip_cursor) }.IpMask.String.as_ptr() as *const i8;
    let mask_bytes = unsafe { CStr::from_ptr(mask_str_ptr).to_bytes() };
    let mask_str = unsafe { from_utf8_unchecked(mask_bytes).to_owned() };
    let mask = mask_str.parse().map_err(|_| ())?;

    let prefix = ip_mask_to_prefix(mask).map_err(|_| ())?;
    IpNetwork::new(ip, prefix).map_err(|_| ())
}
<commit_msg>Update winpcap.rs<commit_after>// Copyright (c) 2014-2016 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 sending and receiving data link layer packets using the WinPcap library

extern crate libc;
extern crate winapi;

use bindings::{bpf, winpcap};
use {DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface};

use ipnetwork::{ip_mask_to_prefix, IpNetwork};

use std::cmp;
use std::collections::VecDeque;
use std::ffi::{CStr, CString};
use std::io;
use std::mem;
use std::slice;
use std::str::from_utf8_unchecked;
use std::sync::Arc;

struct WinPcapAdapter {
    adapter: winpcap::LPADAPTER,
}

impl Drop for WinPcapAdapter {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketCloseAdapter(self.adapter);
        }
    }
}

struct WinPcapPacket {
    packet: winpcap::LPPACKET,
}

impl Drop for WinPcapPacket {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketFreePacket(self.packet);
        }
    }
}

/// WinPcap specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,
}

impl<'a> From<&'a super::Config> for Config {
    fn from(config: &super::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
        }
    }
}

/// Create a datalink channel using the WinPcap library
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<super::Channel> {
    let mut read_buffer = Vec::new();
    read_buffer.resize(config.read_buffer_size, 0u8);

    let mut write_buffer = Vec::new();
    write_buffer.resize(config.write_buffer_size, 0u8);

    let adapter = unsafe {
        let net_if_str = CString::new(network_interface.name.as_bytes()).unwrap();
        winpcap::PacketOpenAdapter(net_if_str.as_ptr() as *mut libc::c_char)
    };
    if adapter.is_null() {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetHwFilter(adapter, winpcap::NDIS_PACKET_TYPE_PROMISCUOUS) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Set kernel buffer size
    let ret = unsafe { winpcap::PacketSetBuff(adapter, config.read_buffer_size as libc::c_int) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Immediate mode
    let ret = unsafe { winpcap::PacketSetMinToCopy(adapter, 1) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let read_packet = unsafe { winpcap::PacketAllocatePacket() };
    if read_packet.is_null() {
        unsafe {
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(read_packet,
                                  read_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.read_buffer_size as winpcap::UINT)
    }

    let write_packet = unsafe { winpcap::PacketAllocatePacket() };
    if write_packet.is_null() {
        unsafe {
            winpcap::PacketFreePacket(read_packet);
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(write_packet,
                                  write_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.write_buffer_size as winpcap::UINT)
    }

    let adapter = Arc::new(WinPcapAdapter { adapter: adapter });
    let sender = Box::new(DataLinkSenderImpl {
        adapter: adapter.clone(),
        _write_buffer: write_buffer,
        packet: WinPcapPacket { packet: write_packet },
    });
    let receiver = Box::new(DataLinkReceiverImpl {
        adapter: adapter,
        _read_buffer: read_buffer,
        packet: WinPcapPacket { packet: read_packet },
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(unsafe { (*read_packet).Length } as usize / 64),
    });
    Ok(super::Channel::Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    adapter: Arc<WinPcapAdapter>,
    _write_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= unsafe { (*self.packet.packet).Length } as usize {
            None
        } else {
            let min = unsafe { cmp::min((*self.packet.packet).Length as usize, len) };
            let slice: &mut [u8] =
                unsafe { slice::from_raw_parts_mut((*self.packet.packet).Buffer as *mut u8, min) };
            for chunk in slice.chunks_mut(packet_size) {
                func(chunk);

                // Make sure the right length of packet is sent
                let old_len = unsafe { (*self.packet.packet).Length };
                unsafe {
                    (*self.packet.packet).Length = packet_size as u32;
                }

                let ret = unsafe {
                    winpcap::PacketSendPacket(self.adapter.adapter, self.packet.packet, 0)
                };

                unsafe {
                    (*self.packet.packet).Length = old_len;
                }

                if ret == 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        self.build_and_send(1,
                            packet.len(),
                            &mut |eh: &mut [u8]| {
                                eh.copy_from_slice(packet);
                            })
    }
}

unsafe impl Send for DataLinkSenderImpl {}
unsafe impl Sync for DataLinkSenderImpl {}

struct DataLinkReceiverImpl {
    adapter: Arc<WinPcapAdapter>,
    _read_buffer: Vec<u8>,
    packet: WinPcapPacket,
    packets: VecDeque<(usize, usize)>,
}

unsafe impl Send for DataLinkReceiverImpl {}
unsafe impl Sync for DataLinkReceiverImpl {}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // NOTE Most of the logic here is identical to FreeBSD/OS X
        while self.packets.is_empty() {
            let ret = unsafe {
                winpcap::PacketReceivePacket(self.adapter.adapter, self.packet.packet, 0)
            };
            let buflen = match ret {
                0 => return Err(io::Error::last_os_error()),
                _ => unsafe { (*self.packet.packet).ulBytesReceived as isize },
            };
            let mut ptr = unsafe { (*self.packet.packet).Buffer };
            let end = unsafe { (*self.packet.packet).Buffer.offset(buflen) };
            while ptr < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                (*self.packet.packet).Buffer as isize;
                    self.packets.push_back((start as usize, (*packet).bh_caplen as usize));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        let slice = unsafe {
            let data = (*self.packet.packet).Buffer as usize + start;
            slice::from_raw_parts(data as *const u8, len)
        };
        Ok(slice)
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use
        // GetAdaptersAddresses
        winpcap::GetAdaptersInfo(&mut tmp, &mut adapters_size);
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    if adapters_size % mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32 != 0 {
        vec_size += 1;
    }

    let mut adapters = Vec::with_capacity(vec_size as usize);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as usize);
    while !cursor.is_null() {
        let mac = unsafe {
            MacAddr((*cursor).Address[0],
                    (*cursor).Address[1],
                    (*cursor).Address[2],
                    (*cursor).Address[3],
                    (*cursor).Address[4],
                    (*cursor).Address[5])
        };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING };
        let mut ips = Vec::new();
        while !ip_cursor.is_null() {
            if let Ok(ip_network) = parse_ip_network(ip_cursor) {
                ips.push(ip_network);
            }
            ip_cursor = unsafe { (*ip_cursor).Next };
        }

        unsafe {
            let name_str_ptr = (*cursor).AdapterName.as_ptr() as *const i8;

            let bytes = CStr::from_ptr(name_str_ptr).to_bytes();
            let name_str = from_utf8_unchecked(bytes).to_owned();

            all_ifaces.push(NetworkInterface {
                name: name_str,
                index: (*cursor).Index,
                mac: Some(mac),
                ips: ips,
                // flags: (*cursor).Type, // FIXME [windows]
                flags: 0,
            });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = vec![0u8; 4096];
    let mut buflen = buf.len() as u32;

    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        buf.resize(buflen as usize, 0);

        // Second call should now work with the correct buffer size. If not, this may be
        // due to some privilege or other unforseen issue.
        if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
            panic!("Unable to get interface list despite increasing buffer size");
        }
    }

    let buf_str = unsafe { from_utf8_unchecked(&buf) };
    let iface_names = buf_str.split("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_owned();
                let next = all_ifaces.iter().filter(|x| name[..].ends_with(&x.name[..])).next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        }
        None => (),
    };

    vec
}

fn parse_ip_network(ip_cursor: winpcap::PIP_ADDR_STRING) -> Result<IpNetwork, ()> {
    let ip_str_ptr = unsafe { &(*ip_cursor) }.IpAddress.String.as_ptr() as *const i8;
    let ip_bytes = unsafe { CStr::from_ptr(ip_str_ptr).to_bytes() };
    let ip_str = unsafe { from_utf8_unchecked(ip_bytes).to_owned() };
    let ip = ip_str.parse().map_err(|_| ())?;

    let mask_str_ptr = unsafe { &(*ip_cursor) }.IpMask.String.as_ptr() as *const i8;
    let mask_bytes = unsafe { CStr::from_ptr(mask_str_ptr).to_bytes() };
    let mask_str = unsafe { from_utf8_unchecked(mask_bytes).to_owned() };
    let mask = mask_str.parse().map_err(|_| ())?;

    let prefix = ip_mask_to_prefix(mask).map_err(|_| ())?;
    IpNetwork::new(ip, prefix).map_err(|_| ())
}
<|endoftext|>


<commit_before>// STD Dependencies -----------------------------------------------------------
use std::fmt;
use std::collections::HashMap;

// Discord Dependencies -------------------------------------------------------
use discord::model::{
    ChannelId, UserId, RoleId, ServerId,
    Member as DiscordMember,
    User as DiscordUser
};

// External Dependencies ------------------------------------------------------
use clock_ticks;

// Internal Dependencies ------------------------------------------------------
use ::bot::BotConfig;

// Member Abstraction ---------------------------------------------------------
#[derive(Debug)]
pub struct Member {
    pub id: UserId,
    pub server_id: ServerId,
    pub name: String,
    pub roles: Vec<RoleId>,
    pub nickname: String,
    pub is_bot: bool,
    pub is_active_bot: bool,
    pub is_admin: bool,
    pub is_uploader: bool,
    pub is_banned: bool,
    pub voice_channel_id: Option<ChannelId>,
    pub last_voice_leave: HashMap<ChannelId, u64>,
    pub mute: bool,
    pub deaf: bool
}

// Public Interface -----------------------------------------------------------
impl Member {

    pub fn from_discord_member(
        discord_member: DiscordMember,
        server_id: ServerId,
        bot_config: &BotConfig

    ) -> Member {
        let mut member = Member::from_discord_user(
            discord_member.user,
            server_id,
            bot_config
        );
        member.roles = discord_member.roles.clone();
        member
    }

    fn from_discord_user(
        user: DiscordUser,
        server_id: ServerId,
        bot_config: &BotConfig

    ) -> Member {
        let nickname = format!("{}#{}", user.name, user.discriminator);
        Member {
            id: user.id,
            server_id: server_id,
            name: user.name.to_string(),
            roles: Vec::new(),
            nickname: nickname.clone(),
            is_bot: user.bot,
            is_active_bot: nickname == bot_config.bot_nickname,
            is_admin: false,
            is_uploader: false,
            is_banned: false,
            voice_channel_id: None,
            last_voice_leave: HashMap::new(),
            mute: false,
            deaf: false
        }
    }

    pub fn should_be_greeted(
        &mut self,
        bot_config: &BotConfig

    ) -> bool {
        // TODO cleanup
        if self.voice_channel_id.is_some() {
            let last_leave = self.last_voice_leave
                                 .entry(self.voice_channel_id.unwrap())
                                 .or_insert(0);

            clock_ticks::precise_time_ms() - *last_leave
                > bot_config.greeting_separation_ms

        } else {
            false
        }
    }

    pub fn left_channel(&mut self, channel_id: &ChannelId) {
        self.last_voice_leave.insert(
            *channel_id,
            clock_ticks::precise_time_ms()
        );
    }

}

// Traits  --------------------------------------------------------------------
impl fmt::Display for Member {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        let kind = match (self.is_bot, self.is_admin, self.is_uploader, self.is_banned) {
            (_, _, _, true) => "<Banned>",
            (true, _, _, false) => "<Bot>",
            (false, true, true, false) => "<Admin, Uploader>",
            (false, true, false, false) => "<Admin>",
            (false, false, true, false) => "<Uploader>",
            (false, false, false, false) => ""
        };

        if self.voice_channel_id.is_some() {
            write!(f, "[Member {}{} (Voice) Mute: {} Deaf: {}]", self.nickname, kind, self.mute, self.deaf)

        } else {
            write!(f, "[Member {}{} (Text)]", self.nickname, kind)
        }

    }
}

<commit_msg>Minor refactoring.<commit_after>// STD Dependencies -----------------------------------------------------------
use std::fmt;
use std::collections::HashMap;

// Discord Dependencies -------------------------------------------------------
use discord::model::{
    ChannelId, UserId, RoleId, ServerId,
    Member as DiscordMember,
    User as DiscordUser
};

// External Dependencies ------------------------------------------------------
use clock_ticks;

// Internal Dependencies ------------------------------------------------------
use ::bot::BotConfig;

// Member Abstraction ---------------------------------------------------------
#[derive(Debug)]
pub struct Member {
    pub id: UserId,
    pub server_id: ServerId,
    pub name: String,
    pub roles: Vec<RoleId>,
    pub nickname: String,
    pub is_bot: bool,
    pub is_active_bot: bool,
    pub is_admin: bool,
    pub is_uploader: bool,
    pub is_banned: bool,
    pub voice_channel_id: Option<ChannelId>,
    pub last_voice_leave: HashMap<ChannelId, u64>,
    pub mute: bool,
    pub deaf: bool
}

// Public Interface -----------------------------------------------------------
impl Member {

    pub fn from_discord_member(
        discord_member: DiscordMember,
        server_id: ServerId,
        bot_config: &BotConfig

    ) -> Member {
        let mut member = Member::from_discord_user(
            discord_member.user,
            server_id,
            bot_config
        );
        member.roles = discord_member.roles.clone();
        member
    }

    fn from_discord_user(
        user: DiscordUser,
        server_id: ServerId,
        bot_config: &BotConfig

    ) -> Member {
        let nickname = format!("{}#{}", user.name, user.discriminator);
        Member {
            id: user.id,
            server_id: server_id,
            name: user.name.to_string(),
            roles: Vec::new(),
            nickname: nickname.clone(),
            is_bot: user.bot,
            is_active_bot: nickname == bot_config.bot_nickname,
            is_admin: false,
            is_uploader: false,
            is_banned: false,
            voice_channel_id: None,
            last_voice_leave: HashMap::new(),
            mute: false,
            deaf: false
        }
    }

    pub fn should_be_greeted(&mut self, bot_config: &BotConfig ) -> bool {
        if self.voice_channel_id.is_some() {
            let now = clock_ticks::precise_time_ms();
            let time = self.last_voice_leave
                           .entry(self.voice_channel_id.unwrap())
                           .or_insert(0);

            now - *time > bot_config.greeting_separation_ms

        } else {
            false
        }
    }

    pub fn left_channel(&mut self, channel_id: &ChannelId) {
        self.last_voice_leave.insert(
            *channel_id,
            clock_ticks::precise_time_ms()
        );
    }

}

// Traits  --------------------------------------------------------------------
impl fmt::Display for Member {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        let kind = match (self.is_bot, self.is_admin, self.is_uploader, self.is_banned) {
            (_, _, _, true) => "<Banned>",
            (true, _, _, false) => "<Bot>",
            (false, true, true, false) => "<Admin, Uploader>",
            (false, true, false, false) => "<Admin>",
            (false, false, true, false) => "<Uploader>",
            (false, false, false, false) => ""
        };

        if self.voice_channel_id.is_some() {
            write!(f, "[Member {}{} (Voice) Mute: {} Deaf: {}]", self.nickname, kind, self.mute, self.deaf)

        } else {
            write!(f, "[Member {}{} (Text)]", self.nickname, kind)
        }

    }
}

<|endoftext|>


<commit_before>use dirs;
use super::error::{S3Error, Result};
use ini::Ini;
use std::collections::HashMap;
use std::env;

/// AWS access credentials: access key, secret key, and optional token.
///
/// # Example
///
/// Loads from the standard AWS credentials file with the given profile name,
/// defaults to "default".
///
/// ```no_run
/// # // Do not execute this as it would cause unit tests to attempt to access
/// # // real user credentials.
/// use s3::credentials::Credentials;
///
/// // Load credentials from `[default]` profile
/// let credentials = Credentials::default();
///
/// // Also loads credentials from `[default]` profile
/// let credentials = Credentials::new(None, None, None, None);
///
/// // Load credentials from `[my-profile]` profile
/// let credentials = Credentials::new(None, None, None, Some("my-profile".into()));
/// ```
///
/// Credentials may also be initialized directly or by the following environment variables:
///
///   - `AWS_ACCESS_KEY_ID`,
///   - `AWS_SECRET_ACCESS_KEY`
///   - `AWS_SESSION_TOKEN`
///
/// The order of preference is arguments, then environment, and finally AWS
/// credentials file.
///
/// ```
/// use s3::credentials::Credentials;
///
/// // Load credentials directly
/// let access_key = String::from("AKIAIOSFODNN7EXAMPLE");
/// let secret_key = String::from("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None);
///
/// // Load credentials from the environment
/// use std::env;
/// env::set_var("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE");
/// env::set_var("AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(None, None, None, None);
/// ```
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Credentials {
    /// AWS public access key.
    pub access_key: String,
    /// AWS secret key.
    pub secret_key: String,
    /// Temporary token issued by AWS service.
    pub token: Option<String>,
    _private: (),
}

impl Credentials {
    /// Initialize Credentials directly with key ID, secret key, and optional
    /// token.
    pub fn new(
        access_key: Option<String>,
        secret_key: Option<String>,
        token: Option<String>,
        profile: Option<String>,
    ) -> Credentials {
        let credentials = match access_key {
            Some(key) => match secret_key {
                Some(secret) => match token {
                    Some(t) => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: Some(t),
                        _private: (),
                    }),
                    None => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: None,
                        _private: (),
                    }),
                },
                None => None,
            },
            None => None,
        };
        match credentials {
            Some(c) => c,
            None => match Credentials::from_env() {
                Ok(c) => c,
                Err(_) => match Credentials::from_profile(profile) {
                    Ok(c) => c,
                    Err(_) => match futures::executor::block_on(Credentials::from_instance_metadata()) {
                        Ok(c) => c,
                        Err(e) => panic!("No credentials provided as arguments, in the environment or in the profile file. \n {}", e)
                    }
                }
            }
        }
    }

    fn from_env() -> Result<Credentials> {
        let access_key = env::var("AWS_ACCESS_KEY_ID")?;
        let secret_key = env::var("AWS_SECRET_ACCESS_KEY")?;
        let token = match env::var("AWS_SESSION_TOKEN") {
            Ok(x) => Some(x),
            Err(_) => None,
        };
        Ok(Credentials {
            access_key,
            secret_key,
            token,
            _private: (),
        })
    }

    async fn from_instance_metadata() -> Result<Credentials> {
        if !Credentials::is_ec2() {
            return Err(S3Error::from("Not an EC2 instance"));
        }

        let resp: HashMap<String, String> =
            reqwest::get("http://169.254.169.254/latest/meta-data/iam/info").await?.json().await?;
        let credentials = if let Some(arn) = resp.get("InstanceProfileArn") {
            if let Some(role) = arn.split('/').last() {
                let resp: HashMap<String, String> = reqwest::get(&format!(
                    "http://169.254.169.254/latest/meta-data/iam/security-credentials/{}",
                    role
                )).await?
                .json().await?;
                let access_key = resp.get("AccessKeyId").unwrap().clone();
                let secret_key = resp.get("SecretAccessKey").unwrap().clone();
                let token = Some(resp.get("Token").unwrap().clone());
                Some(Credentials {
                    access_key,
                    secret_key,
                    token,
                    _private: (),
                })
            } else {
                None
            }
        } else {
            None
        };

        Ok(credentials.unwrap())
    }

    fn is_ec2() -> bool {
        if let Ok(uuid) = std::fs::read_to_string("/sys/hypervisor/uuid") {
            if uuid.len() >= 3 && &uuid[..3] == "ec2" {
                return true;
            }
        }
        if let Ok(uuid) = std::fs::read_to_string("/sys/class/dmi/id/board_vendor") {
            if uuid.len() >= 10 && &uuid[..10] == "Amazon EC2" {
                return true;
            }
        }
        false
    }

    pub fn from_profile(section: Option<String>) -> Result<Credentials> {
        let home_dir = match dirs::home_dir() {
            Some(path) => Ok(path),
            None => Err(S3Error::from("Invalid home dir")),
        };
        let profile = format!("{}/.aws/credentials", home_dir?.display());
        let conf = Ini::load_from_file(&profile)?;
        let section = match section {
            Some(s) => s,
            None => String::from("default"),
        };
        let mut access_key = Err(S3Error::from("Missing aws_access_key_id section"));
        let mut secret_key = Err(S3Error::from("Missing aws_secret_access_key section"));
        let mut token = None;
        if let Some(data) = conf.section(Some(section)) {
            access_key = match data.get("aws_access_key_id") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_access_key_id section")),
            };
            secret_key = match data.get("aws_secret_access_key") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_secret_access_key section")),
            };
            token = match data.get("aws_security_token") {
                Some(x) => Some(x.to_owned()),
                None => None,
            }
        }

        Ok(Credentials {
            access_key: access_key?,
            secret_key: secret_key?,
            token,
            _private: (),
        })
    }
}

impl Default for Credentials {
    fn default() -> Self {
        Credentials::new(None, None, None, None)
    }
}
<commit_msg>Fargate support (#75)<commit_after>use dirs;
use super::error::{S3Error, Result};
use ini::Ini;
use std::collections::HashMap;
use std::env;

/// AWS access credentials: access key, secret key, and optional token.
///
/// # Example
///
/// Loads from the standard AWS credentials file with the given profile name,
/// defaults to "default".
///
/// ```no_run
/// # // Do not execute this as it would cause unit tests to attempt to access
/// # // real user credentials.
/// use s3::credentials::Credentials;
///
/// // Load credentials from `[default]` profile
/// let credentials = Credentials::default();
///
/// // Also loads credentials from `[default]` profile
/// let credentials = Credentials::new(None, None, None, None);
///
/// // Load credentials from `[my-profile]` profile
/// let credentials = Credentials::new(None, None, None, Some("my-profile".into()));
/// ```
///
/// Credentials may also be initialized directly or by the following environment variables:
///
///   - `AWS_ACCESS_KEY_ID`,
///   - `AWS_SECRET_ACCESS_KEY`
///   - `AWS_SESSION_TOKEN`
///
/// The order of preference is arguments, then environment, and finally AWS
/// credentials file.
///
/// ```
/// use s3::credentials::Credentials;
///
/// // Load credentials directly
/// let access_key = String::from("AKIAIOSFODNN7EXAMPLE");
/// let secret_key = String::from("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None);
///
/// // Load credentials from the environment
/// use std::env;
/// env::set_var("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE");
/// env::set_var("AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(None, None, None, None);
/// ```
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Credentials {
    /// AWS public access key.
    pub access_key: String,
    /// AWS secret key.
    pub secret_key: String,
    /// Temporary token issued by AWS service.
    pub token: Option<String>,
    _private: (),
}

impl Credentials {
    /// Initialize Credentials directly with key ID, secret key, and optional
    /// token.
    pub fn new(
        access_key: Option<String>,
        secret_key: Option<String>,
        token: Option<String>,
        profile: Option<String>,
    ) -> Credentials {
        let credentials = match access_key {
            Some(key) => match secret_key {
                Some(secret) => match token {
                    Some(t) => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: Some(t),
                        _private: (),
                    }),
                    None => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: None,
                        _private: (),
                    }),
                },
                None => None,
            },
            None => None,
        };
        match credentials {
            Some(c) => c,
            None => match Credentials::from_env() {
                Ok(c) => c,
                Err(_) => match Credentials::from_profile(profile) {
                    Ok(c) => c,
                    Err(_) => match futures::executor::block_on(Credentials::from_instance_metadata()) {
                        Ok(c) => c,
                        Err(e) => panic!("No credentials provided as arguments, in the environment or in the profile file. \n {}", e)
                    }
                }
            }
        }
    }

    fn from_env() -> Result<Credentials> {
        let access_key = env::var("AWS_ACCESS_KEY_ID")?;
        let secret_key = env::var("AWS_SECRET_ACCESS_KEY")?;
        let token = match env::var("AWS_SESSION_TOKEN") {
            Ok(x) => Some(x),
            Err(_) => None,
        };
        Ok(Credentials {
            access_key,
            secret_key,
            token,
            _private: (),
        })
    }

    async fn from_instance_metadata() -> Result<Credentials> {
        if !Credentials::is_ec2() {
            return Err(S3Error::from("Not an EC2 instance"));
        }
        let resp:HashMap<String, String> = match env::var("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI") {
            Ok(credentials_path) => {
                Some(reqwest::get(&format!("http://169.254.170.2{}",credentials_path)).await?.json().await?)
            },
            Err(_) => {
                let resp: HashMap<String, String> =
                    reqwest::get("http://169.254.169.254/latest/meta-data/iam/info").await?.json().await?;
                if let Some(arn) = resp.get("InstanceProfileArn") {
                    if let Some(role) = arn.split('/').last() {
                        Some(reqwest::get(&format!(
                            "http://169.254.169.254/latest/meta-data/iam/security-credentials/{}",
                            role
                        )).await?
                        .json().await?)
                    } else {
                        None
                    }
                } else {
                    None
                }
            }
        }.unwrap();

        let access_key = resp.get("AccessKeyId").unwrap().clone();
        let secret_key = resp.get("SecretAccessKey").unwrap().clone();
        let token = Some(resp.get("Token").unwrap().clone());
        Ok(Credentials {
            access_key,
            secret_key,
            token,
            _private: (),
        })

    }

    fn is_ec2() -> bool {
        if let Ok(uuid) = std::fs::read_to_string("/sys/hypervisor/uuid") {
            if uuid.len() >= 3 && &uuid[..3] == "ec2" {
                return true;
            }
        }
        if let Ok(uuid) = std::fs::read_to_string("/sys/class/dmi/id/board_vendor") {
            if uuid.len() >= 10 && &uuid[..10] == "Amazon EC2" {
                return true;
            }
        }
        false
    }

    pub fn from_profile(section: Option<String>) -> Result<Credentials> {
        let home_dir = match dirs::home_dir() {
            Some(path) => Ok(path),
            None => Err(S3Error::from("Invalid home dir")),
        };
        let profile = format!("{}/.aws/credentials", home_dir?.display());
        let conf = Ini::load_from_file(&profile)?;
        let section = match section {
            Some(s) => s,
            None => String::from("default"),
        };
        let mut access_key = Err(S3Error::from("Missing aws_access_key_id section"));
        let mut secret_key = Err(S3Error::from("Missing aws_secret_access_key section"));
        let mut token = None;
        if let Some(data) = conf.section(Some(section)) {
            access_key = match data.get("aws_access_key_id") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_access_key_id section")),
            };
            secret_key = match data.get("aws_secret_access_key") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_secret_access_key section")),
            };
            token = match data.get("aws_security_token") {
                Some(x) => Some(x.to_owned()),
                None => None,
            }
        }

        Ok(Credentials {
            access_key: access_key?,
            secret_key: secret_key?,
            token,
            _private: (),
        })
    }
}

impl Default for Credentials {
    fn default() -> Self {
        Credentials::new(None, None, None, None)
    }
}
<|endoftext|>


<commit_before>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::MetaDataConsumer;

// Will return true when the unwrapped value of `$option` and `$compare`
// match or `$option` is `Option::None`, otherwise false.
macro_rules! optional_eq (
  ($compare: expr, $option: expr) => (
    $option.map_or(true, |compare| $compare == compare);
  );
);

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_cue_sheet("path/to/file.flac") {
///   Ok(cue_sheet) => {
///     // Use the cue_sheet variable...
///   }
///   Err(error)    => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors
///
/// ```no_run
/// use flac::metadata;
///
/// let cue_sheet = metadata::get_cue_sheet("path/to/file.flac").unwrap();
/// ```
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let result = metadata::get_picture("path/to/file.flac",
///                                    Some(PictureType::FileIcon),
///                                    Some("image/gif"),
///                                    None, None, None,
///                                    None, None);
///
/// match result {
///   Ok(picture) => {
///     // Use the picture variable...
///   }
///   Err(error)  => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let picture = metadata::get_picture("path/to/file.flac",
///                                     Some(PictureType::FileIcon),
///                                     Some("image/gif"),
///                                     None, None, None,
///                                     None, None).unwrap();
/// ```
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::io::ErrorKind;

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<commit_msg>Change `optional_eq!` macro to a function<commit_after>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::MetaDataConsumer;

// Will return true when the unwrapped value of `$option` and `$compare`
// match or `$option` is `Option::None`, otherwise false.
#[inline]
pub fn optional_eq<T: Eq>(option: Option<T>, other: T) -> bool {
  option.map_or(true, |value| value == other)
}

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_cue_sheet("path/to/file.flac") {
///   Ok(cue_sheet) => {
///     // Use the cue_sheet variable...
///   }
///   Err(error)    => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors
///
/// ```no_run
/// use flac::metadata;
///
/// let cue_sheet = metadata::get_cue_sheet("path/to/file.flac").unwrap();
/// ```
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let result = metadata::get_picture("path/to/file.flac",
///                                    Some(PictureType::FileIcon),
///                                    Some("image/gif"),
///                                    None, None, None,
///                                    None, None);
///
/// match result {
///   Ok(picture) => {
///     // Use the picture variable...
///   }
///   Err(error)  => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
/// use flac::metadata::PictureType;
///
/// let picture = metadata::get_picture("path/to/file.flac",
///                                     Some(PictureType::FileIcon),
///                                     Some("image/gif"),
///                                     None, None, None,
///                                     None, None).unwrap();
/// ```
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::io::ErrorKind;

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<|endoftext|>


<commit_before>//! This module provides a way to construct a `File`.
//! It is intended to be completely decoupled from the
//! parser, so as to allow to evolve the tree representation
//! and the parser algorithm independently.
//!
//! The `Sink` trait is the bridge between the parser and the
//! tree builder: the parser produces a stream of events like
//! `start node`, `finish node`, and `FileBuilder` converts
//! this stream to a real tree.
use {
    lexer::Token,
    parser_impl::Sink,
    SyntaxKind::{self, TOMBSTONE},
};

/// `Parser` produces a flat list of `Event`s.
/// They are converted to a tree-structure in
/// a separate pass, via `TreeBuilder`.
#[derive(Debug)]
pub(crate) enum Event {
    /// This event signifies the start of the node.
    /// It should be either abandoned (in which case the
    /// `kind` is `TOMBSTONE`, and the event is ignored),
    /// or completed via a `Finish` event.
    ///
    /// All tokens between a `Start` and a `Finish` would
    /// become the children of the respective node.
    ///
    /// For left-recursive syntactic constructs, the parser produces
    /// a child node before it sees a parent. `forward_parent`
    /// exists to allow to tweak parent-child relationships.
    ///
    /// Consider this path
    ///
    /// foo::bar
    ///
    /// The events for it would look like this:
    ///
    ///
    /// START(PATH) IDENT('foo') FINISH START(PATH) COLONCOLON IDENT('bar') FINISH
    ///       |                          /\
    ///       |                          |
    ///       +------forward-parent------+
    ///
    /// And the tree would look like this
    ///
    ///    +--PATH---------+
    ///    |   |           |
    ///    |   |           |
    ///    |  '::'       'bar'
    ///    |
    ///   PATH
    ///    |
    ///   'foo'
    ///
    /// See also `CompletedMarker::precede`.
    Start {
        kind: SyntaxKind,
        forward_parent: Option<u32>,
    },

    /// Complete the previous `Start` event
    Finish,

    /// Produce a single leaf-element.
    /// `n_raw_tokens` is used to glue complex contextual tokens.
    /// For example, lexer tokenizes `>>` as `>`, `>`, and
    /// `n_raw_tokens = 2` is used to produced a single `>>`.
    Token {
        kind: SyntaxKind,
        n_raw_tokens: u8,
    },

    Error {
        msg: String,
    },
}

pub(super) fn process<'a, S: Sink<'a>>(builder: &mut S, tokens: &[Token], events: Vec<Event>) {
    let mut next_tok_idx = 0;
    let eat_ws = |idx: &mut usize, builder: &mut S| {
        while let Some(token) = tokens.get(*idx) {
            if !token.kind.is_trivia() {
                break;
            }
            builder.leaf(token.kind, token.len);
            *idx += 1
        }
    };

    let mut depth = 0;
    let mut holes = Vec::new();
    let mut forward_parents = Vec::new();

    for (i, event) in events.iter().enumerate() {
        if holes.last() == Some(&i) {
            holes.pop();
            continue;
        }

        match event {
            &Event::Start {
                kind: TOMBSTONE, ..
            } => (),

            &Event::Start { .. } => {
                forward_parents.clear();
                let mut idx = i;
                loop {
                    let (kind, fwd) = match events[idx] {
                        Event::Start {
                            kind,
                            forward_parent,
                        } => (kind, forward_parent),
                        _ => unreachable!(),
                    };
                    forward_parents.push((idx, kind));
                    if let Some(fwd) = fwd {
                        idx += fwd as usize;
                    } else {
                        break;
                    }
                }
                for &(idx, kind) in forward_parents.iter().into_iter().rev() {
                    if depth > 0 {
                        eat_ws(&mut next_tok_idx, builder);
                    }
                    depth += 1;
                    builder.start_internal(kind);
                    holes.push(idx);
                }
                holes.pop();
            }
            &Event::Finish => {
                depth -= 1;
                if depth == 0 {
                    eat_ws(&mut next_tok_idx, builder);
                }

                builder.finish_internal();
            }
            &Event::Token {
                kind,
                mut n_raw_tokens,
            } => {
                eat_ws(&mut next_tok_idx, builder);
                let mut len = 0.into();
                for _ in 0..n_raw_tokens {
                    len += tokens[next_tok_idx].len;
                    next_tok_idx += 1;
                }
                builder.leaf(kind, len);
            }
            &Event::Error { ref msg } => builder.error(msg.clone()),
        }
    }
}
<commit_msg>Simplify event processing<commit_after>//! This module provides a way to construct a `File`.
//! It is intended to be completely decoupled from the
//! parser, so as to allow to evolve the tree representation
//! and the parser algorithm independently.
//!
//! The `Sink` trait is the bridge between the parser and the
//! tree builder: the parser produces a stream of events like
//! `start node`, `finish node`, and `FileBuilder` converts
//! this stream to a real tree.
use std::mem;
use {
    lexer::Token,
    parser_impl::Sink,
    SyntaxKind::{self, TOMBSTONE},
};

/// `Parser` produces a flat list of `Event`s.
/// They are converted to a tree-structure in
/// a separate pass, via `TreeBuilder`.
#[derive(Debug)]
pub(crate) enum Event {
    /// This event signifies the start of the node.
    /// It should be either abandoned (in which case the
    /// `kind` is `TOMBSTONE`, and the event is ignored),
    /// or completed via a `Finish` event.
    ///
    /// All tokens between a `Start` and a `Finish` would
    /// become the children of the respective node.
    ///
    /// For left-recursive syntactic constructs, the parser produces
    /// a child node before it sees a parent. `forward_parent`
    /// exists to allow to tweak parent-child relationships.
    ///
    /// Consider this path
    ///
    /// foo::bar
    ///
    /// The events for it would look like this:
    ///
    ///
    /// START(PATH) IDENT('foo') FINISH START(PATH) COLONCOLON IDENT('bar') FINISH
    ///       |                          /\
    ///       |                          |
    ///       +------forward-parent------+
    ///
    /// And the tree would look like this
    ///
    ///    +--PATH---------+
    ///    |   |           |
    ///    |   |           |
    ///    |  '::'       'bar'
    ///    |
    ///   PATH
    ///    |
    ///   'foo'
    ///
    /// See also `CompletedMarker::precede`.
    Start {
        kind: SyntaxKind,
        forward_parent: Option<u32>,
    },

    /// Complete the previous `Start` event
    Finish,

    /// Produce a single leaf-element.
    /// `n_raw_tokens` is used to glue complex contextual tokens.
    /// For example, lexer tokenizes `>>` as `>`, `>`, and
    /// `n_raw_tokens = 2` is used to produced a single `>>`.
    Token {
        kind: SyntaxKind,
        n_raw_tokens: u8,
    },

    Error {
        msg: String,
    },
}

pub(super) fn process<'a, S: Sink<'a>>(builder: &mut S, tokens: &[Token], mut events: Vec<Event>) {
    fn tombstone() -> Event {
        Event::Start { kind: TOMBSTONE, forward_parent: None }
    }
    let eat_ws = |idx: &mut usize, builder: &mut S| {
        while let Some(token) = tokens.get(*idx) {
            if !token.kind.is_trivia() {
                break;
            }
            builder.leaf(token.kind, token.len);
            *idx += 1
        }
    };

    let events: &mut [Event] = &mut events;
    let mut depth = 0;
    let mut forward_parents = Vec::new();
    let mut next_tok_idx = 0;
    for i in 0..events.len() {
        match mem::replace(&mut events[i], tombstone()) {
            Event::Start {
                kind: TOMBSTONE, ..
            } => (),

            Event::Start { kind, forward_parent } => {
                forward_parents.push(kind);
                let mut idx = i;
                let mut fp = forward_parent;
                while let Some(fwd) = fp {
                    idx += fwd as usize;
                    fp = match mem::replace(&mut events[idx], tombstone()) {
                        Event::Start {
                            kind,
                            forward_parent,
                        } => {
                            forward_parents.push(kind);
                            forward_parent
                        },
                        _ => unreachable!(),
                    };
                }
                for kind in forward_parents.drain(..).rev() {
                    if depth > 0 {
                        eat_ws(&mut next_tok_idx, builder);
                    }
                    depth += 1;
                    builder.start_internal(kind);
                }
            }
            Event::Finish => {
                depth -= 1;
                if depth == 0 {
                    eat_ws(&mut next_tok_idx, builder);
                }

                builder.finish_internal();
            }
            Event::Token {
                kind,
                mut n_raw_tokens,
            } => {
                eat_ws(&mut next_tok_idx, builder);
                let mut len = 0.into();
                for _ in 0..n_raw_tokens {
                    len += tokens[next_tok_idx].len;
                    next_tok_idx += 1;
                }
                builder.leaf(kind, len);
            }
            Event::Error { msg } => builder.error(msg),
        }
    }
}
<|endoftext|>