<commit_before>//! Unix handling of child processes
//!
//! Right now the only "fancy" thing about this is how we implement the
//! `Future` implementation on `Child` to get the exit status. Unix offers
//! no way to register a child with epoll, and the only real way to get a
//! notification when a process exits is the SIGCHLD signal.
//!
//! Signal handling in general is *super* hairy and complicated, and it's even
//! more complicated here with the fact that signals are coalesced, so we may
//! not get a SIGCHLD-per-child.
//!
//! Our best approximation here is to check *all spawned processes* for all
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
//! the `tokio-signal` crate, which is a stream over signals being received.
//!
//! Later when we poll the process's exit status we simply check to see if a
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
//! keep trying.
//!
//! Note that this means that this isn't really scalable, but then again
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
//! bad in theory...

extern crate libc;
extern crate tokio_signal;

use std::io;
use std::os::unix::prelude::*;
use std::process::{self, ExitStatus};

use futures::future::FlattenStream;
use futures::{Future, Poll, Async, Stream};
use mio::unix::{EventedFd, UnixReady};
use mio::{PollOpt, Ready, Token};
use mio::event::Evented;
use mio;
use self::libc::c_int;
use self::tokio_signal::unix::Signal;
use std::fmt;
use tokio_io::IoFuture;
use tokio_core::reactor::{Handle, PollEvented};

#[must_use = "futures do nothing unless polled"]
pub struct Child {
    inner: process::Child,
    reaped: bool,
    sigchld: FlattenStream<IoFuture<Signal>>,
}

impl fmt::Debug for Child {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Child")
            .field("pid", &self.inner.id())
            .field("inner", &self.inner)
            .field("reaped", &self.reaped)
            .field("sigchld", &"..")
            .finish()
    }
}

impl Child {
    pub fn new(inner: process::Child, handle: &Handle) -> Child {
        Child {
            inner: inner,
            reaped: false,
            sigchld: Signal::new(libc::SIGCHLD, handle).flatten_stream(),
        }
    }

    pub fn register_stdin(&mut self, handle: &Handle)
                          -> io::Result<Option<ChildStdin>> {
        stdio(self.inner.stdin.take(), handle)
    }

    pub fn register_stdout(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStdout>> {
        stdio(self.inner.stdout.take(), handle)
    }

    pub fn register_stderr(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStderr>> {
        stdio(self.inner.stderr.take(), handle)
    }

    pub fn id(&self) -> u32 {
        self.inner.id()
    }

    pub fn kill(&mut self) -> io::Result<()> {
        if self.reaped {
            Ok(())
        } else {
            self.inner.kill()?;
            let mut status = 0;
            let id = self.id() as c_int;
            unsafe { libc::waitpid(id, &mut status, 0) };
            Ok(())
        }
    }

    pub fn poll_exit(&mut self) -> Poll<ExitStatus, io::Error> {
        assert!(!self.reaped);
        loop {
            // Ensure that once we've successfully waited we won't try to
            // `kill` above.
            if let Some(e) = try!(self.try_wait()) {
                self.reaped = true;
                return Ok(e.into())
            }

            // If the child hasn't exited yet, then it's our responsibility to
            // ensure the current task gets notified when it might be able to
            // make progress.
            //
            // As described in `spawn` above, we just indicate that we can
            // next make progress once a SIGCHLD is received.
            if try!(self.sigchld.poll()).is_not_ready() {
                return Ok(Async::NotReady)
            }
        }
    }

    fn try_wait(&self) -> io::Result<Option<ExitStatus>> {
        let id = self.id() as c_int;
        let mut status = 0;
        loop {
            match unsafe { libc::waitpid(id, &mut status, libc::WNOHANG) } {
                0 => return Ok(None),
                n if n < 0 => {
                    let err = io::Error::last_os_error();
                    if err.kind() == io::ErrorKind::Interrupted {
                        continue
                    }
                    return Err(err)
                }
                n => {
                    assert_eq!(n, id);
                    return Ok(Some(ExitStatus::from_raw(status)))
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct Fd<T>(T);

impl<T: io::Read> io::Read for Fd<T> {
    fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
        self.0.read(bytes)
    }
}

impl<T: io::Write> io::Write for Fd<T> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0.write(bytes)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

pub type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
pub type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
pub type ChildStderr = PollEvented<Fd<process::ChildStderr>>;

impl<T> Evented for Fd<T> where T: AsRawFd {
    fn register(&self,
                poll: &mio::Poll,
                token: Token,
                interest: Ready,
                opts: PollOpt)
                -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).register(poll,
                                                token,
                                                interest | UnixReady::hup(),
                                                opts)
    }

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

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

fn stdio<T>(option: Option<T>, handle: &Handle)
            -> io::Result<Option<PollEvented<Fd<T>>>>
    where T: AsRawFd
{
    let io = match option {
        Some(io) => io,
        None => return Ok(None),
    };

    // Set the fd to nonblocking before we pass it to the event loop
    unsafe {
        let fd = io.as_raw_fd();
        let r = libc::fcntl(fd, libc::F_GETFL);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
        let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
    }
    let io = try!(PollEvented::new(Fd(io), handle));
    Ok(Some(io))
}
<commit_msg>process: Unix: mark child as reaped on kill<commit_after>//! Unix handling of child processes
//!
//! Right now the only "fancy" thing about this is how we implement the
//! `Future` implementation on `Child` to get the exit status. Unix offers
//! no way to register a child with epoll, and the only real way to get a
//! notification when a process exits is the SIGCHLD signal.
//!
//! Signal handling in general is *super* hairy and complicated, and it's even
//! more complicated here with the fact that signals are coalesced, so we may
//! not get a SIGCHLD-per-child.
//!
//! Our best approximation here is to check *all spawned processes* for all
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
//! the `tokio-signal` crate, which is a stream over signals being received.
//!
//! Later when we poll the process's exit status we simply check to see if a
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
//! keep trying.
//!
//! Note that this means that this isn't really scalable, but then again
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
//! bad in theory...

extern crate libc;
extern crate tokio_signal;

use std::io;
use std::os::unix::prelude::*;
use std::process::{self, ExitStatus};

use futures::future::FlattenStream;
use futures::{Future, Poll, Async, Stream};
use mio::unix::{EventedFd, UnixReady};
use mio::{PollOpt, Ready, Token};
use mio::event::Evented;
use mio;
use self::tokio_signal::unix::Signal;
use std::fmt;
use tokio_io::IoFuture;
use tokio_core::reactor::{Handle, PollEvented};

#[must_use = "futures do nothing unless polled"]
pub struct Child {
    inner: process::Child,
    reaped: bool,
    sigchld: FlattenStream<IoFuture<Signal>>,
}

impl fmt::Debug for Child {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Child")
            .field("pid", &self.inner.id())
            .field("inner", &self.inner)
            .field("reaped", &self.reaped)
            .field("sigchld", &"..")
            .finish()
    }
}

impl Child {
    pub fn new(inner: process::Child, handle: &Handle) -> Child {
        Child {
            inner: inner,
            reaped: false,
            sigchld: Signal::new(libc::SIGCHLD, handle).flatten_stream(),
        }
    }

    pub fn register_stdin(&mut self, handle: &Handle)
                          -> io::Result<Option<ChildStdin>> {
        stdio(self.inner.stdin.take(), handle)
    }

    pub fn register_stdout(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStdout>> {
        stdio(self.inner.stdout.take(), handle)
    }

    pub fn register_stderr(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStderr>> {
        stdio(self.inner.stderr.take(), handle)
    }

    pub fn id(&self) -> u32 {
        self.inner.id()
    }

    pub fn kill(&mut self) -> io::Result<()> {
        if !self.reaped {
            // NB: SIGKILL cannnot be caught, so the process will definitely exit immediately.
            // We're not waiting for the process itself but for the kernel to execute the kill.
            self.inner.kill()?;
            let _ = self.try_wait(true);
        }

        Ok(())
    }

    pub fn poll_exit(&mut self) -> Poll<ExitStatus, io::Error> {
        loop {
            // Ensure that once we've successfully waited we won't try to
            // `kill` above.
            if let Some(e) = try!(self.try_wait(false)) {
                return Ok(e.into())
            }

            // If the child hasn't exited yet, then it's our responsibility to
            // ensure the current task gets notified when it might be able to
            // make progress.
            //
            // As described in `spawn` above, we just indicate that we can
            // next make progress once a SIGCHLD is received.
            if try!(self.sigchld.poll()).is_not_ready() {
                return Ok(Async::NotReady)
            }
        }
    }

    fn try_wait(&mut self, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
        assert!(!self.reaped);
        let exit = try!(try_wait_process(self.id() as libc::pid_t, block_on_wait));

        if let Some(_) = exit {
            self.reaped = true;
        }

        Ok(exit)
    }
}

fn try_wait_process(id: libc::pid_t, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
    let wait_flags = if block_on_wait { 0 } else { libc::WNOHANG };
    let mut status = 0;

    loop {
        match unsafe { libc::waitpid(id, &mut status, wait_flags) } {
            0 => return Ok(None),
            n if n < 0 => {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                }
                return Err(err)
            }
            n => {
                assert_eq!(n, id);
                return Ok(Some(ExitStatus::from_raw(status)))
            }
        }
    }
}

#[derive(Debug)]
pub struct Fd<T>(T);

impl<T: io::Read> io::Read for Fd<T> {
    fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
        self.0.read(bytes)
    }
}

impl<T: io::Write> io::Write for Fd<T> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0.write(bytes)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

pub type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
pub type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
pub type ChildStderr = PollEvented<Fd<process::ChildStderr>>;

impl<T> Evented for Fd<T> where T: AsRawFd {
    fn register(&self,
                poll: &mio::Poll,
                token: Token,
                interest: Ready,
                opts: PollOpt)
                -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).register(poll,
                                                token,
                                                interest | UnixReady::hup(),
                                                opts)
    }

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

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

fn stdio<T>(option: Option<T>, handle: &Handle)
            -> io::Result<Option<PollEvented<Fd<T>>>>
    where T: AsRawFd
{
    let io = match option {
        Some(io) => io,
        None => return Ok(None),
    };

    // Set the fd to nonblocking before we pass it to the event loop
    unsafe {
        let fd = io.as_raw_fd();
        let r = libc::fcntl(fd, libc::F_GETFL);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
        let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
    }
    let io = try!(PollEvented::new(Fd(io), handle));
    Ok(Some(io))
}
<|endoftext|>


<commit_before>use urlparse::{Query, parse_qs, urlparse, urlunparse};

/// Re-encode a query string for Rust
fn encode_querystring(query: Query) -> Option<String> {
  let mut query_components: Vec<String> = vec![];
  for (key, value) in query {
    for v in value.iter() {
      query_components.push(format!("{}={}", key, v));
    }
  }
  if query_components.len() > 0 {
    Some(query_components.join("&"))
  } else {
    None
  }
}

/// Strip tracking junk and URL suffixes.
pub fn tidy_url(url: &str) -> String {
    let mut parsed_url = urlparse(url);

    // Always get the desktop version of Twitter URLs
    if parsed_url.netloc == "mobile.twitter.com" {
        parsed_url.netloc = String::from("twitter.com");
    }

    // Remove any tracking junk from Amazon URLs so they're not a
    // ridiculous length
    if parsed_url.netloc == "www.amazon.co.uk" {
        parsed_url.query = None;
        let mut new_path = String::from("");
        for component in parsed_url.path.split("ref=") {
            new_path = String::from(component);
            break;
        }
        parsed_url.path = new_path;
    }

    // Strip tracking junk from Medium, Mashable and Buzzfeed
    if (parsed_url.netloc == "medium.com") ||
       (parsed_url.netloc == "www.buzzfeed.com") ||
       (parsed_url.netloc == "mashable.com") {
            parsed_url.fragment = None;
    }

    // Remove '#notes' from Tumblr URLs
    if parsed_url.netloc.ends_with("tumblr.com") {
        let new_fragment: Option<String> = match parsed_url.fragment {
            Some(fragment) => {
                if fragment == "notes" {
                    None
                } else {
                    Some(fragment)
                }
            }
            None => None,
        };
        parsed_url.fragment = new_fragment;
    }

  // Remove &feature=youtu.be from YouTube URLs
  if parsed_url.netloc.ends_with("youtube.com") {
    parsed_url.query = match parsed_url.query {
      Some(qs) => {
        let mut query = parse_qs(&qs);
        query.remove("feature");
        encode_querystring(query)
      },
      None => None,
    };
  }

  // Remove any UTM tracking parameters from URLs
  parsed_url.query = match parsed_url.query {
    Some(qs) => {
      let mut query = parse_qs(&qs);
      let utm_keys: Vec<_> = query
        .keys()
        .filter(|key| key.starts_with("utm_"))
        .map(|k| k.clone())
        .collect();
      for key in utm_keys {
        query.remove(&key);
      }
      encode_querystring(query)
    },
    None => None
  };

  urlunparse(parsed_url)
}

macro_rules! tidy_url_tests {
    ($($name:ident: $value:expr,)*) => {
    $(
        #[test]
        fn $name() {
            let (input, expected) = $value;
            assert_eq!(expected, tidy_url(input));
        }
    )*
    }
}

tidy_url_tests! {
    url_tw_0:  ("https://mobile.twitter.com/Breaking911/status/822589852596191235",
                "https://twitter.com/Breaking911/status/822589852596191235"),
    url_tw_1:  ("https://twitter.com/elise3aum/status/822903824268533762",
                "https://twitter.com/elise3aum/status/822903824268533762"),

    url_az_0:  ("https://www.amazon.co.uk/dp/B01DFKBL68/ref=gw_aucc_comb_AB_clean_2?pf_rd_r=N1MF2PWADRHS55427ETG&pf_rd_p=c11a2c11-c670-46ff-87a1-c1eef4fad652",
                "https://www.amazon.co.uk/dp/B01DFKBL68/"),

    url_md_0:  ("https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20#.sjyskxdsz",
                "https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20"),

    url_tm_0:  ("http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539#notes",
                "http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539"),

    url_8:   ("https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man#.biqmkzz8kn",
              "https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man"),
    url_9:   ("http://mashable.com/2016/03/21/apple-liam-recycling-robot/#b9y4lv3m4qqX",
              "http://mashable.com/2016/03/21/apple-liam-recycling-robot/"),

  youtube_plain_video: (
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
  ),

  youtube_with_timestamp: (
    "https://www.youtube.com/watch?v=D68cUzqcTrg&feature=youtu.be#t=1m",
    "https://www.youtube.com/watch?v=D68cUzqcTrg#t=1m",
  ),

  youtube_with_trailing_feature: (
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0&feature=youtu.be",
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0",
  ),

  youtube_with_leading_feature: (
    "https://www.youtube.com/watch?feature=youtu.be&v=oPo4n9tBPsk",
    "https://www.youtube.com/watch?v=oPo4n9tBPsk",
  ),

  youtube_with_only_feature: (
    "https://www.youtube.com/watch?feature=youtu.be",
    "https://www.youtube.com/watch",
  ),

  youtube_channel: (
    "https://www.youtube.com/user/TheQIElves",
    "https://www.youtube.com/user/TheQIElves"
  ),

  single_utm_tracker: (
    "https://example.com?utm_medium=social",
    "https://example.com",
  ),

  multiple_utm_tracker: (
    "https://example.com?utm_medium=social&utm_source=twitter",
    "https://example.com",
  ),

  multiple_utm_tracker_with_others: (
    "https://example.com?utm_medium=social&foo=bar&utm_source=twitter",
    "https://example.com?foo=bar",
  ),
}
<commit_msg>Add a failing test for URL encoding bugs<commit_after>use urlparse::{Query, parse_qs, urlparse, urlunparse};

/// Re-encode a query string for Rust
fn encode_querystring(query: Query) -> Option<String> {
  let mut query_components: Vec<String> = vec![];
  for (key, value) in query {
    for v in value.iter() {
      query_components.push(format!("{}={}", key, v));
    }
  }
  if query_components.len() > 0 {
    Some(query_components.join("&"))
  } else {
    None
  }
}

/// Strip tracking junk and URL suffixes.
pub fn tidy_url(url: &str) -> String {
    let mut parsed_url = urlparse(url);

    // Always get the desktop version of Twitter URLs
    if parsed_url.netloc == "mobile.twitter.com" {
        parsed_url.netloc = String::from("twitter.com");
    }

    // Remove any tracking junk from Amazon URLs so they're not a
    // ridiculous length
    if parsed_url.netloc == "www.amazon.co.uk" {
        parsed_url.query = None;
        let mut new_path = String::from("");
        for component in parsed_url.path.split("ref=") {
            new_path = String::from(component);
            break;
        }
        parsed_url.path = new_path;
    }

    // Strip tracking junk from Medium, Mashable and Buzzfeed
    if (parsed_url.netloc == "medium.com") ||
       (parsed_url.netloc == "www.buzzfeed.com") ||
       (parsed_url.netloc == "mashable.com") {
            parsed_url.fragment = None;
    }

    // Remove '#notes' from Tumblr URLs
    if parsed_url.netloc.ends_with("tumblr.com") {
        let new_fragment: Option<String> = match parsed_url.fragment {
            Some(fragment) => {
                if fragment == "notes" {
                    None
                } else {
                    Some(fragment)
                }
            }
            None => None,
        };
        parsed_url.fragment = new_fragment;
    }

  // Remove &feature=youtu.be from YouTube URLs
  if parsed_url.netloc.ends_with("youtube.com") {
    parsed_url.query = match parsed_url.query {
      Some(qs) => {
        let mut query = parse_qs(&qs);
        query.remove("feature");
        encode_querystring(query)
      },
      None => None,
    };
  }

  // Remove any UTM tracking parameters from URLs
  parsed_url.query = match parsed_url.query {
    Some(qs) => {
      let mut query = parse_qs(&qs);
      let utm_keys: Vec<_> = query
        .keys()
        .filter(|key| key.starts_with("utm_"))
        .map(|k| k.clone())
        .collect();
      for key in utm_keys {
        query.remove(&key);
      }
      encode_querystring(query)
    },
    None => None
  };

  urlunparse(parsed_url)
}

macro_rules! tidy_url_tests {
    ($($name:ident: $value:expr,)*) => {
    $(
        #[test]
        fn $name() {
            let (input, expected) = $value;
            assert_eq!(expected, tidy_url(input));
        }
    )*
    }
}

tidy_url_tests! {
    url_tw_0:  ("https://mobile.twitter.com/Breaking911/status/822589852596191235",
                "https://twitter.com/Breaking911/status/822589852596191235"),
    url_tw_1:  ("https://twitter.com/elise3aum/status/822903824268533762",
                "https://twitter.com/elise3aum/status/822903824268533762"),

    url_az_0:  ("https://www.amazon.co.uk/dp/B01DFKBL68/ref=gw_aucc_comb_AB_clean_2?pf_rd_r=N1MF2PWADRHS55427ETG&pf_rd_p=c11a2c11-c670-46ff-87a1-c1eef4fad652",
                "https://www.amazon.co.uk/dp/B01DFKBL68/"),

    url_md_0:  ("https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20#.sjyskxdsz",
                "https://medium.com/@anildash/forget-why-its-time-to-get-to-work-c49ac5f0da20"),

    url_tm_0:  ("http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539#notes",
                "http://azurelunatic.tumblr.com/post/155525051123/things-about-hufflepuffs-539"),

    url_8:   ("https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man#.biqmkzz8kn",
              "https://www.buzzfeed.com/katienotopoulos/the-end-of-apple-man"),
    url_9:   ("http://mashable.com/2016/03/21/apple-liam-recycling-robot/#b9y4lv3m4qqX",
              "http://mashable.com/2016/03/21/apple-liam-recycling-robot/"),

  youtube_plain_video: (
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
    "https://www.youtube.com/watch?v=zB4I68XVPzQ",
  ),

  youtube_with_timestamp: (
    "https://www.youtube.com/watch?v=D68cUzqcTrg&feature=youtu.be#t=1m",
    "https://www.youtube.com/watch?v=D68cUzqcTrg#t=1m",
  ),

  youtube_with_trailing_feature: (
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0&feature=youtu.be",
    "https://www.youtube.com/watch?v=PbJqNa0_Oz0",
  ),

  youtube_with_leading_feature: (
    "https://www.youtube.com/watch?feature=youtu.be&v=oPo4n9tBPsk",
    "https://www.youtube.com/watch?v=oPo4n9tBPsk",
  ),

  youtube_with_only_feature: (
    "https://www.youtube.com/watch?feature=youtu.be",
    "https://www.youtube.com/watch",
  ),

  youtube_channel: (
    "https://www.youtube.com/user/TheQIElves",
    "https://www.youtube.com/user/TheQIElves"
  ),

  single_utm_tracker: (
    "https://example.com?utm_medium=social",
    "https://example.com",
  ),

  multiple_utm_tracker: (
    "https://example.com?utm_medium=social&utm_source=twitter",
    "https://example.com",
  ),

  multiple_utm_tracker_with_others: (
    "https://example.com?utm_medium=social&foo=bar&utm_source=twitter",
    "https://example.com?foo=bar",
  ),

  url_with_spaces: (
    "https://example.com?foo=bar%20baz",
    "https://example.com?foo=bar%20baz"
  ),

  url_with_numerals: (
    "https://example.com?foo=bar0baz",
    "https://example.com?foo=bar0baz"
  ),
}
<|endoftext|>


<commit_before>use std::io::Read;
use errors::Result;

pub trait ZeroCompress : Read {
    fn decode_vint64(&mut self) -> Result<i64> {
        let mut raw_buffer = vec![0u8; 1];
        try!(self.read_exact(&mut raw_buffer));

        let value = raw_buffer[0] as i8;
        let len = if value >= -112 {
            1
        } else if value < -120 {
            -119 - value
        } else {
            -111 - value
        };

        let mut val = 0i64;
        if len == 1 {
            val = value as i64;
        } else {
            for _ in 0..(len - 1) {
                try!(self.read_exact(&mut raw_buffer));
                val = val << 8;
                val = val | (raw_buffer[0] as i64 & 0xFF)
            }
        }

        if value < -120 || (value >= -112 && value < 0) {
            Ok(val ^ -1)
        } else {
            Ok(val)
        }
    }
}

impl<R> ZeroCompress for R where R: Read
{}
<commit_msg>util: test coverage around vint decoding<commit_after>use std::io::Read;
use errors::Result;

pub trait ZeroCompress : Read {
    fn decode_vint64(&mut self) -> Result<i64> {
        let mut raw_buffer = vec![0u8; 1];
        try!(self.read_exact(&mut raw_buffer));

        let value = raw_buffer[0] as i8;
        let len = if value >= -112 {
            1
        } else if value < -120 {
            // [2, 8], negative result
            -119 - value
        } else {
            // [2, 8], positive result
            -111 - value
        };

        let mut val = 0i64;
        if len == 1 {
            val = value as i64;
        } else {
            // shifts buffer to make room for new set-o-bits, ors in the new byte
            for _ in 0..(len - 1) {
                try!(self.read_exact(&mut raw_buffer));
                val = val << 8;
                val = val | (raw_buffer[0] as i64 & 0xFF)
            }
        }

        if value < -120 || (value >= -112 && value < 0) {
            Ok(val ^ -1)
        } else {
            Ok(val)
        }
    }
}

impl<R> ZeroCompress for R where R: Read
{}

#[cfg(test)]
mod tests {
    use super::ZeroCompress;
    use std::io::{Read, Cursor};

    #[test]
    fn decodes_single_byte() {
        let mut buf = Cursor::new(vec![0b0111_1111]);

        assert_eq!(127, buf.decode_vint64().unwrap());
    }

    #[test]
    fn decodes_multi_byte() {
        let mut buf = Cursor::new(vec![0b1000_1101, 0b1000_0000, 0b0100_0000, 0b0010_1101]);

        assert_eq!(8_405_037, buf.decode_vint64().unwrap());
    }

    #[test]
    fn decodes_multi_byte_negative() {
        let mut buf = Cursor::new(vec![0b1000_0101, 0b1000_0000, 0b0100_0000, 0b0010_1101]);

        assert_eq!(-8_405_038, buf.decode_vint64().unwrap());
    }
}
<|endoftext|>


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

//! Miscellaneous utilities for low level networking

extern crate libc;

use packet::PrimitiveValues;
use packet::ip::IpNextHeaderProtocol;
use datalink::NetworkInterface;
use pnet_macros_support::types::u16be;

use std::fmt;
use std::str::FromStr;
use std::u8;
use std::net::{IpAddr, Ipv4Addr};
use std::mem;
use std::slice;

use internal;
use sockets;

/// A MAC address
#[derive(PartialEq, Eq, Clone, Copy, Hash)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);

impl MacAddr {
    /// Construct a new MacAddr
    pub fn new(a: u8, b: u8, c: u8, d: u8, e: u8, f: u8) -> MacAddr {
        MacAddr(a, b, c, d, e, f)
    }
}

impl PrimitiveValues for MacAddr {
    type T = (u8, u8, u8, u8, u8, u8);
    fn to_primitive_values(&self) -> (u8, u8, u8, u8, u8, u8) {
        (self.0, self.1, self.2, self.3, self.4, self.5)
    }
}

impl fmt::Display for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt,
               "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}",
               self.0,
               self.1,
               self.2,
               self.3,
               self.4,
               self.5)
    }
}

impl fmt::Debug for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(self, fmt)
    }
}

// FIXME Is this the right way to do this? Which occurs is an implementation
//       issue rather than actually defined - is it useful to provide these
//       errors, or would it be better to just give ()?
/// Represents an error which occurred whilst parsing a MAC address
#[derive(Copy, Debug, PartialEq, Eq, Clone)]
pub enum ParseMacAddrErr {
    /// The MAC address has too many components, eg. 00:11:22:33:44:55:66
    TooManyComponents,
    /// The MAC address has too few components, eg. 00:11
    TooFewComponents,
    /// One of the components contains an invalid value, eg. 00:GG:22:33:44:55
    InvalidComponent,
}

impl FromStr for MacAddr {
    type Err = ParseMacAddrErr;
    fn from_str(s: &str) -> Result<MacAddr, ParseMacAddrErr> {
        let mut parts = [0u8; 6];
        let splits = s.split(':');
        let mut i = 0;
        for split in splits {
            if i == 6 {
                return Err(ParseMacAddrErr::TooManyComponents);
            }
            match u8::from_str_radix(split, 16) {
                Ok(b) if split.len() != 0 => parts[i] = b,
                _ => return Err(ParseMacAddrErr::InvalidComponent),
            }
            i += 1;
        }

        if i == 6 {
            Ok(MacAddr(parts[0], parts[1], parts[2], parts[3], parts[4], parts[5]))
        } else {
            Err(ParseMacAddrErr::TooFewComponents)
        }
    }
}

#[test]
fn mac_addr_from_str() {
    assert_eq!("00:00:00:00:00:00".parse(), Ok(MacAddr(0, 0, 0, 0, 0, 0)));
    assert_eq!("ff:ff:ff:ff:ff:ff".parse(),
               Ok(MacAddr(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)));
    assert_eq!("12:34:56:78:90:ab".parse(),
               Ok(MacAddr(0x12, 0x34, 0x56, 0x78, 0x90, 0xAB)));
    assert_eq!("::::::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("0::::::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("::::0::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooFewComponents));
    assert_eq!("12:34:56:78:".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78:90".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooFewComponents));
    assert_eq!("12:34:56:78:90:".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78:90:00:00".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooManyComponents));
    assert_eq!("xx:xx:xx:xx:xx:xx".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const sockets::SockAddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0],
                              (*sll).sll_addr[1],
                              (*sll).sll_addr[2],
                              (*sll).sll_addr[3],
                              (*sll).sll_addr[4],
                              (*sll).sll_addr[5]);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<sockets::SockAddrStorage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const sockets::SockAddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr((*sdl).sdl_data[nlen] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<sockets::SockAddrStorage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

/// Get a list of available network interfaces for the current machine.
/// Deprecated. Instead use the implementation available for your backend.
/// The default one is at `pnet::datalink::interfaces`.
#[deprecated(note="Moved to datalink::interfaces()")]
#[inline]
pub fn get_network_interfaces() -> Vec<NetworkInterface> {
    ::datalink::interfaces()
}

/// Convert value to byte array
pub trait Octets {
    /// Output type - bytes array
    type Output;

    /// Return value as bytes (big-endian order)
    fn octets(&self) -> Self::Output;
}

impl Octets for u64 {
    type Output = [u8; 8];

    fn octets(&self) -> Self::Output {
        [(*self >> 56) as u8, (*self >> 48) as u8, (*self >> 40) as u8, (*self >> 32) as u8,
         (*self >> 24) as u8, (*self >> 16) as u8, (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u32 {
    type Output = [u8; 4];

    fn octets(&self) -> Self::Output {
        [(*self >> 24) as u8, (*self >> 16) as u8 , (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u16 {
    type Output = [u8; 2];

    fn octets(&self) -> Self::Output {
        [(*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u8 {
    type Output = [u8; 1];

    fn octets(&self) -> Self::Output {
        [*self]
    }
}

/// Calculates a checksum. Used by ipv4 and icmp. The two bytes starting at `skipword * 2` will be
/// ignored. Supposed to be the checksum field, which is regarded as zero during calculation.
/// To not skip any word, set `skipword` to `data.len()`.
pub fn checksum(data: &[u8], skipword: Option<usize>) -> u16be {
    let mut sum = sum_be_words(data, skipword);
    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }
    !sum as u16
}

/// Calculate the checksum for a packet built on IPv4
pub fn ipv4_checksum(data: &[u8],
                     ipv4_source: Ipv4Addr,
                     ipv4_destination: Ipv4Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    // IPv4 source
    let octets = ipv4_source.octets();
    sum += (octets[0] as u32) << 8 | octets[1] as u32;
    sum += (octets[2] as u32) << 8 | octets[3] as u32;

    // IPv4 destination
    let octets = ipv4_destination.octets();
    sum += (octets[0] as u32) << 8 | octets[1] as u32;
    sum += (octets[2] as u32) << 8 | octets[3] as u32;

    // IPv4 Next level protocol
    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    // Length
    sum += data.len() as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, None);

    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }

    !sum as u16
}

/// Sum all words (16 bit chunks) in the given data. The word at word offset `skipword` will be
/// ignored if passed. Each word is treated as big endian.
fn sum_be_words(data: &[u8], skipword: Option<usize>) -> u32 {
    let len = data.len();
    let wdata: &[u16] = unsafe { slice::from_raw_parts(data.as_ptr() as *const u16, len / 2) };

    let mut sum = 0u32;
    let mut i = 0;
    if let Some(skipword) = skipword {
        assert!(skipword <= wdata.len());
        while i < skipword {
            sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
            i += 1;
        }
        i += 1;
    }
    while i < wdata.len() {
        sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
        i += 1;
    }
    // If the length is odd, make sure to checksum the final byte
    if len & 1 != 0 {
        sum += (data[len - 1] as u32) << 8;
    }

    sum
}

#[cfg(all(test, feature = "benchmark"))]
mod checksum_benchmarks {
    use super::checksum;
    use test::{black_box, Bencher};

    #[bench]
    fn bench_checksum_small(b: &mut Bencher) {
        let data = vec![99u8; 20];
        b.iter(|| {
            checksum(black_box(&data), Some(5))
        });
    }

    #[bench]
    fn bench_checksum_large(b: &mut Bencher) {
        let data = vec![123u8; 1024];
        b.iter(|| {
            checksum(black_box(&data), Some(5))
        });
    }
}
<commit_msg>rustfmt benchmarks<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Miscellaneous utilities for low level networking

extern crate libc;

use packet::PrimitiveValues;
use packet::ip::IpNextHeaderProtocol;
use datalink::NetworkInterface;
use pnet_macros_support::types::u16be;

use std::fmt;
use std::str::FromStr;
use std::u8;
use std::net::{IpAddr, Ipv4Addr};
use std::mem;
use std::slice;

use internal;
use sockets;

/// A MAC address
#[derive(PartialEq, Eq, Clone, Copy, Hash)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);

impl MacAddr {
    /// Construct a new MacAddr
    pub fn new(a: u8, b: u8, c: u8, d: u8, e: u8, f: u8) -> MacAddr {
        MacAddr(a, b, c, d, e, f)
    }
}

impl PrimitiveValues for MacAddr {
    type T = (u8, u8, u8, u8, u8, u8);
    fn to_primitive_values(&self) -> (u8, u8, u8, u8, u8, u8) {
        (self.0, self.1, self.2, self.3, self.4, self.5)
    }
}

impl fmt::Display for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt,
               "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}",
               self.0,
               self.1,
               self.2,
               self.3,
               self.4,
               self.5)
    }
}

impl fmt::Debug for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(self, fmt)
    }
}

// FIXME Is this the right way to do this? Which occurs is an implementation
//       issue rather than actually defined - is it useful to provide these
//       errors, or would it be better to just give ()?
/// Represents an error which occurred whilst parsing a MAC address
#[derive(Copy, Debug, PartialEq, Eq, Clone)]
pub enum ParseMacAddrErr {
    /// The MAC address has too many components, eg. 00:11:22:33:44:55:66
    TooManyComponents,
    /// The MAC address has too few components, eg. 00:11
    TooFewComponents,
    /// One of the components contains an invalid value, eg. 00:GG:22:33:44:55
    InvalidComponent,
}

impl FromStr for MacAddr {
    type Err = ParseMacAddrErr;
    fn from_str(s: &str) -> Result<MacAddr, ParseMacAddrErr> {
        let mut parts = [0u8; 6];
        let splits = s.split(':');
        let mut i = 0;
        for split in splits {
            if i == 6 {
                return Err(ParseMacAddrErr::TooManyComponents);
            }
            match u8::from_str_radix(split, 16) {
                Ok(b) if split.len() != 0 => parts[i] = b,
                _ => return Err(ParseMacAddrErr::InvalidComponent),
            }
            i += 1;
        }

        if i == 6 {
            Ok(MacAddr(parts[0], parts[1], parts[2], parts[3], parts[4], parts[5]))
        } else {
            Err(ParseMacAddrErr::TooFewComponents)
        }
    }
}

#[test]
fn mac_addr_from_str() {
    assert_eq!("00:00:00:00:00:00".parse(), Ok(MacAddr(0, 0, 0, 0, 0, 0)));
    assert_eq!("ff:ff:ff:ff:ff:ff".parse(),
               Ok(MacAddr(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)));
    assert_eq!("12:34:56:78:90:ab".parse(),
               Ok(MacAddr(0x12, 0x34, 0x56, 0x78, 0x90, 0xAB)));
    assert_eq!("::::::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("0::::::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("::::0::".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooFewComponents));
    assert_eq!("12:34:56:78:".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78:90".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooFewComponents));
    assert_eq!("12:34:56:78:90:".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
    assert_eq!("12:34:56:78:90:00:00".parse::<MacAddr>(),
               Err(ParseMacAddrErr::TooManyComponents));
    assert_eq!("xx:xx:xx:xx:xx:xx".parse::<MacAddr>(),
               Err(ParseMacAddrErr::InvalidComponent));
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const sockets::SockAddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0],
                              (*sll).sll_addr[1],
                              (*sll).sll_addr[2],
                              (*sll).sll_addr[3],
                              (*sll).sll_addr[4],
                              (*sll).sll_addr[5]);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<sockets::SockAddrStorage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const sockets::SockAddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr((*sdl).sdl_data[nlen] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8);

            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                                  mem::size_of::<sockets::SockAddrStorage>());

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

/// Get a list of available network interfaces for the current machine.
/// Deprecated. Instead use the implementation available for your backend.
/// The default one is at `pnet::datalink::interfaces`.
#[deprecated(note="Moved to datalink::interfaces()")]
#[inline]
pub fn get_network_interfaces() -> Vec<NetworkInterface> {
    ::datalink::interfaces()
}

/// Convert value to byte array
pub trait Octets {
    /// Output type - bytes array
    type Output;

    /// Return value as bytes (big-endian order)
    fn octets(&self) -> Self::Output;
}

impl Octets for u64 {
    type Output = [u8; 8];

    fn octets(&self) -> Self::Output {
        [(*self >> 56) as u8, (*self >> 48) as u8, (*self >> 40) as u8, (*self >> 32) as u8,
         (*self >> 24) as u8, (*self >> 16) as u8, (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u32 {
    type Output = [u8; 4];

    fn octets(&self) -> Self::Output {
        [(*self >> 24) as u8, (*self >> 16) as u8 , (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u16 {
    type Output = [u8; 2];

    fn octets(&self) -> Self::Output {
        [(*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u8 {
    type Output = [u8; 1];

    fn octets(&self) -> Self::Output {
        [*self]
    }
}

/// Calculates a checksum. Used by ipv4 and icmp. The two bytes starting at `skipword * 2` will be
/// ignored. Supposed to be the checksum field, which is regarded as zero during calculation.
/// To not skip any word, set `skipword` to `data.len()`.
pub fn checksum(data: &[u8], skipword: Option<usize>) -> u16be {
    let mut sum = sum_be_words(data, skipword);
    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }
    !sum as u16
}

/// Calculate the checksum for a packet built on IPv4
pub fn ipv4_checksum(data: &[u8],
                     ipv4_source: Ipv4Addr,
                     ipv4_destination: Ipv4Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    // IPv4 source
    let octets = ipv4_source.octets();
    sum += (octets[0] as u32) << 8 | octets[1] as u32;
    sum += (octets[2] as u32) << 8 | octets[3] as u32;

    // IPv4 destination
    let octets = ipv4_destination.octets();
    sum += (octets[0] as u32) << 8 | octets[1] as u32;
    sum += (octets[2] as u32) << 8 | octets[3] as u32;

    // IPv4 Next level protocol
    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    // Length
    sum += data.len() as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, None);

    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }

    !sum as u16
}

/// Sum all words (16 bit chunks) in the given data. The word at word offset `skipword` will be
/// ignored if passed. Each word is treated as big endian.
fn sum_be_words(data: &[u8], skipword: Option<usize>) -> u32 {
    let len = data.len();
    let wdata: &[u16] = unsafe { slice::from_raw_parts(data.as_ptr() as *const u16, len / 2) };

    let mut sum = 0u32;
    let mut i = 0;
    if let Some(skipword) = skipword {
        assert!(skipword <= wdata.len());
        while i < skipword {
            sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
            i += 1;
        }
        i += 1;
    }
    while i < wdata.len() {
        sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
        i += 1;
    }
    // If the length is odd, make sure to checksum the final byte
    if len & 1 != 0 {
        sum += (data[len - 1] as u32) << 8;
    }

    sum
}

#[cfg(all(test, feature = "benchmark"))]
mod checksum_benchmarks {
    use super::checksum;
    use test::{Bencher, black_box};

    #[bench]
    fn bench_checksum_small(b: &mut Bencher) {
        let data = vec![99u8; 20];
        b.iter(|| checksum(black_box(&data), Some(5)));
    }

    #[bench]
    fn bench_checksum_large(b: &mut Bencher) {
        let data = vec![123u8; 1024];
        b.iter(|| checksum(black_box(&data), Some(5)));
    }
}
<|endoftext|>


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

use std::fs;
use std::path::Path;
use walkdir::WalkDir;

// Needed to set the script mode to executable.
#[cfg(unix)]
use std::os::unix::fs::OpenOptionsExt;
// FIXME: what about Windows?  Are default ACLs executable?

use errors::*;

/// Convert a `&Path` to a UTF-8 `&str`
pub fn path_to_str(path: &Path) -> Result<&str> {
    path.to_str().ok_or_else(|| {
        ErrorKind::Msg(format!("path is not valid UTF-8 '{}'", path.display())).into()
    })
}

/// Wrap `fs::copy` with a nicer error message
pub fn copy<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> Result<u64> {
    fs::copy(&from, &to)
        .chain_err(|| format!("failed to copy '{}' to '{}'",
                              from.as_ref().display(), to.as_ref().display()))
}

/// Wrap `fs::create_dir` with a nicer error message
pub fn create_dir<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::create_dir(&path)
        .chain_err(|| format!("failed to create dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::create_dir_all` with a nicer error message
pub fn create_dir_all<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::create_dir_all(&path)
        .chain_err(|| format!("failed to create dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::OpenOptions::create_new().open()` as executable, with a nicer error message
pub fn create_new_executable<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    let mut options = fs::OpenOptions::new();
    options.write(true).create_new(true);
    #[cfg(unix)] options.mode(0o755);
    options.open(&path)
        .chain_err(|| format!("failed to create file '{}'", path.as_ref().display()))
}

/// Wrap `fs::OpenOptions::create_new().open()`, with a nicer error message
pub fn create_new_file<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    fs::OpenOptions::new().write(true).create_new(true).open(&path)
        .chain_err(|| format!("failed to create file '{}'", path.as_ref().display()))
}

/// Wrap `fs::File::open()` with a nicer error message
pub fn open_file<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    fs::File::open(&path)
        .chain_err(|| format!("failed to open file '{}'", path.as_ref().display()))
}

/// Wrap `remove_dir_all` with a nicer error message
pub fn remove_dir_all<P: AsRef<Path>>(path: P) -> Result<()> {
    ::remove_dir_all::remove_dir_all(path.as_ref())
        .chain_err(|| format!("failed to remove dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::remove_file` with a nicer error message
pub fn remove_file<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::remove_file(path.as_ref())
        .chain_err(|| format!("failed to remove file '{}'", path.as_ref().display()))
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.
pub fn copy_recursive(src: &Path, dst: &Path) -> Result<()> {
    copy_with_callback(src, dst, |_, _| Ok(()))
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.  Invokes a callback for each path visited.
pub fn copy_with_callback<F>(src: &Path, dst: &Path, mut callback: F) -> Result<()>
    where F: FnMut(&Path, fs::FileType) -> Result<()>
{
    for entry in WalkDir::new(src).min_depth(1) {
        let entry = entry?;
        let file_type = entry.file_type();
        let path = entry.path().strip_prefix(src)?;
        let dst = dst.join(path);

        if file_type.is_dir() {
            create_dir(&dst)?;
        } else {
            copy(entry.path(), dst)?;
        }
        callback(&path, file_type)?;
    }
    Ok(())
}

/// Create an "actor" with default values and setters for all fields.
macro_rules! actor {
    ($( #[ $attr:meta ] )+ pub struct $name:ident {
        $( $( #[ $field_attr:meta ] )+ $field:ident : $type:ty = $default:expr, )*
    }) => {
        $( #[ $attr ] )+
        pub struct $name {
            $( $( #[ $field_attr ] )+ $field : $type, )*
        }

        impl Default for $name {
            fn default() -> Self {
                $name {
                    $( $field : $default.into(), )*
                }
            }
        }

        impl $name {
            $( $( #[ $field_attr ] )+
            pub fn $field<T: Into<$type>>(&mut self, value: T) -> &mut Self {
                self.$field = value.into();
                self
            })+
        }
    }
}
<commit_msg>Copy symlinks<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::fs;
use std::path::Path;
use walkdir::WalkDir;

// Needed to set the script mode to executable.
#[cfg(unix)]
use std::os::unix::fs::OpenOptionsExt;
// FIXME: what about Windows?  Are default ACLs executable?

#[cfg(unix)]
use std::os::unix::fs::symlink as symlink_file;
#[cfg(windows)]
use std::os::windows::fs::symlink_file;

use errors::*;

/// Convert a `&Path` to a UTF-8 `&str`
pub fn path_to_str(path: &Path) -> Result<&str> {
    path.to_str().ok_or_else(|| {
        ErrorKind::Msg(format!("path is not valid UTF-8 '{}'", path.display())).into()
    })
}

/// Wrap `fs::copy` with a nicer error message
pub fn copy<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> Result<u64> {
    if fs::symlink_metadata(&from)?.file_type().is_symlink() {
        let link = fs::read_link(&from)?;
        symlink_file(link, &to)?;
        Ok(0)
    } else {
        fs::copy(&from, &to)
            .chain_err(|| format!("failed to copy '{}' to '{}'",
                                  from.as_ref().display(), to.as_ref().display()))
    }
}

/// Wrap `fs::create_dir` with a nicer error message
pub fn create_dir<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::create_dir(&path)
        .chain_err(|| format!("failed to create dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::create_dir_all` with a nicer error message
pub fn create_dir_all<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::create_dir_all(&path)
        .chain_err(|| format!("failed to create dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::OpenOptions::create_new().open()` as executable, with a nicer error message
pub fn create_new_executable<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    let mut options = fs::OpenOptions::new();
    options.write(true).create_new(true);
    #[cfg(unix)] options.mode(0o755);
    options.open(&path)
        .chain_err(|| format!("failed to create file '{}'", path.as_ref().display()))
}

/// Wrap `fs::OpenOptions::create_new().open()`, with a nicer error message
pub fn create_new_file<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    fs::OpenOptions::new().write(true).create_new(true).open(&path)
        .chain_err(|| format!("failed to create file '{}'", path.as_ref().display()))
}

/// Wrap `fs::File::open()` with a nicer error message
pub fn open_file<P: AsRef<Path>>(path: P) -> Result<fs::File> {
    fs::File::open(&path)
        .chain_err(|| format!("failed to open file '{}'", path.as_ref().display()))
}

/// Wrap `remove_dir_all` with a nicer error message
pub fn remove_dir_all<P: AsRef<Path>>(path: P) -> Result<()> {
    ::remove_dir_all::remove_dir_all(path.as_ref())
        .chain_err(|| format!("failed to remove dir '{}'", path.as_ref().display()))
}

/// Wrap `fs::remove_file` with a nicer error message
pub fn remove_file<P: AsRef<Path>>(path: P) -> Result<()> {
    fs::remove_file(path.as_ref())
        .chain_err(|| format!("failed to remove file '{}'", path.as_ref().display()))
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.
pub fn copy_recursive(src: &Path, dst: &Path) -> Result<()> {
    copy_with_callback(src, dst, |_, _| Ok(()))
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.  Invokes a callback for each path visited.
pub fn copy_with_callback<F>(src: &Path, dst: &Path, mut callback: F) -> Result<()>
    where F: FnMut(&Path, fs::FileType) -> Result<()>
{
    for entry in WalkDir::new(src).min_depth(1) {
        let entry = entry?;
        let file_type = entry.file_type();
        let path = entry.path().strip_prefix(src)?;
        let dst = dst.join(path);

        if file_type.is_dir() {
            create_dir(&dst)?;
        } else {
            copy(entry.path(), dst)?;
        }
        callback(&path, file_type)?;
    }
    Ok(())
}

/// Create an "actor" with default values and setters for all fields.
macro_rules! actor {
    ($( #[ $attr:meta ] )+ pub struct $name:ident {
        $( $( #[ $field_attr:meta ] )+ $field:ident : $type:ty = $default:expr, )*
    }) => {
        $( #[ $attr ] )+
        pub struct $name {
            $( $( #[ $field_attr ] )+ $field : $type, )*
        }

        impl Default for $name {
            fn default() -> Self {
                $name {
                    $( $field : $default.into(), )*
                }
            }
        }

        impl $name {
            $( $( #[ $field_attr ] )+
            pub fn $field<T: Into<$type>>(&mut self, value: T) -> &mut Self {
                self.$field = value.into();
                self
            })+
        }
    }
}
<|endoftext|>


<commit_before>use geometry::Point;
use models::{Particle, Vector};

/// Optimized version of `Vec::retain`
///
/// We achieve better performance by renouncing to keep the original order of the `Vec`
pub fn fast_retain<T, F>(vec: &mut Vec<T>, mut f: F)
where F: FnMut(&T) -> bool {
    let mut i = 0;
    while i < vec.len() {
        if !f(&vec[i]) {
            vec.swap_remove(i);
        }

        i += 1;
    }
}

/// Generates a new explosion of the given intensity at the given position.
/// This works best with values between 5 and 25
pub fn make_explosion(particles: &mut Vec<Particle>, position: &Point, intensity: u8) {
    use itertools_num;
    for rotation in itertools_num::linspace(0.0, 2.0 * ::std::f64::consts::PI, 30) {
        for ttl in (1..intensity).map(|x| (x as f64) / 10.0) {
            particles.push(Particle::new(Vector::new(position.clone(), rotation), ttl));
        }
    }
}
<commit_msg>Fix util::fast_retain<commit_after>use geometry::Point;
use models::{Particle, Vector};

/// Optimized version of `Vec::retain`
///
/// We achieve better performance by renouncing to keep the original order of the `Vec`
pub fn fast_retain<T, F>(vec: &mut Vec<T>, mut f: F)
where F: FnMut(&T) -> bool {
    let mut i = 0;
    while i < vec.len() {
        if !f(&vec[i]) {
            vec.swap_remove(i);
        } else {
            i += 1;
        }
    }
}

/// Generates a new explosion of the given intensity at the given position.
/// This works best with values between 5 and 25
pub fn make_explosion(particles: &mut Vec<Particle>, position: &Point, intensity: u8) {
    use itertools_num;
    for rotation in itertools_num::linspace(0.0, 2.0 * ::std::f64::consts::PI, 30) {
        for ttl in (1..intensity).map(|x| (x as f64) / 10.0) {
            particles.push(Particle::new(Vector::new(position.clone(), rotation), ttl));
        }
    }
}

#[test]
fn test_fast_retain() {
    let mut xs = vec![42; 100];
    fast_retain(&mut xs, |&x| x < 42);
    assert_eq!(xs.len(), 0);
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 fd developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use crate::exec;
use crate::exit_codes::ExitCode;
use crate::fshelper;
use crate::internal::{opts::FdOptions, osstr_to_bytes, MAX_BUFFER_LENGTH};
use crate::output;

use std::borrow::Cow;
use std::error::Error;
use std::ffi::OsStr;
use std::io;
use std::path::PathBuf;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use ignore::overrides::OverrideBuilder;
use ignore::{self, WalkBuilder};
use regex::bytes::Regex;

/// The receiver thread can either be buffering results or directly streaming to the console.
enum ReceiverMode {
    /// Receiver is still buffering in order to sort the results, if the search finishes fast
    /// enough.
    Buffering,

    /// Receiver is directly printing results to the output.
    Streaming,
}

/// The Worker threads can result in a valid entry having PathBuf or an error.
pub enum WorkerResult {
    Entry(PathBuf),
    Error(ignore::Error),
}

/// Recursively scan the given search path for files / pathnames matching the pattern.
///
/// If the `--exec` argument was supplied, this will create a thread pool for executing
/// jobs in parallel from a given command line and the discovered paths. Otherwise, each
/// path will simply be written to standard output.
pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<FdOptions>) -> ExitCode {
    let mut path_iter = path_vec.iter();
    let first_path_buf = path_iter
        .next()
        .expect("Error: Path vector can not be empty");
    let (tx, rx) = channel();

    let mut override_builder = OverrideBuilder::new(first_path_buf.as_path());

    for pattern in &config.exclude_patterns {
        let res = override_builder.add(pattern);
        if res.is_err() {
            print_error_and_exit!("Malformed exclude pattern '{}'", pattern);
        }
    }
    let overrides = override_builder.build().unwrap_or_else(|_| {
        print_error_and_exit!("Mismatch in exclude patterns");
    });

    let mut walker = WalkBuilder::new(first_path_buf.as_path());
    walker
        .hidden(config.ignore_hidden)
        .ignore(config.read_fdignore)
        .parents(config.read_fdignore || config.read_vcsignore)
        .git_ignore(config.read_vcsignore)
        .git_global(config.read_vcsignore)
        .git_exclude(config.read_vcsignore)
        .overrides(overrides)
        .follow_links(config.follow_links)
        .max_depth(config.max_depth);

    if config.read_fdignore {
        walker.add_custom_ignore_filename(".fdignore");
    }

    for ignore_file in &config.ignore_files {
        let result = walker.add_ignore(ignore_file);
        match result {
            Some(ignore::Error::Partial(_)) => (),
            Some(err) => {
                print_error!(
                    "{}",
                    format!(
                        "Malformed pattern in custom ignore file '{}': {}.",
                        ignore_file.to_string_lossy(),
                        err.description()
                    )
                );
            }
            None => (),
        }
    }

    for path_entry in path_iter {
        walker.add(path_entry.as_path());
    }

    let parallel_walker = walker.threads(config.threads).build_parallel();

    let wants_to_quit = Arc::new(AtomicBool::new(false));
    if config.ls_colors.is_some() && config.command.is_none() {
        let wq = Arc::clone(&wants_to_quit);
        ctrlc::set_handler(move || {
            if wq.load(Ordering::Relaxed) {
                // Ctrl-C has been pressed twice, exit NOW
                process::exit(ExitCode::KilledBySigint.into());
            } else {
                wq.store(true, Ordering::Relaxed);
            }
        })
        .unwrap();
    }

    // Spawn the thread that receives all results through the channel.
    let receiver_thread = spawn_receiver(&config, &wants_to_quit, rx);

    // Spawn the sender threads.
    spawn_senders(&config, &wants_to_quit, pattern, parallel_walker, tx);

    // Wait for the receiver thread to print out all results.
    let exit_code = receiver_thread.join().unwrap();

    if wants_to_quit.load(Ordering::Relaxed) {
        process::exit(ExitCode::KilledBySigint.into());
    }

    exit_code
}

fn spawn_receiver(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    rx: Receiver<WorkerResult>,
) -> thread::JoinHandle<ExitCode> {
    let config = Arc::clone(config);
    let wants_to_quit = Arc::clone(wants_to_quit);

    let show_filesystem_errors = config.show_filesystem_errors;
    let threads = config.threads;

    thread::spawn(move || {
        // This will be set to `Some` if the `--exec` argument was supplied.
        if let Some(ref cmd) = config.command {
            if cmd.in_batch_mode() {
                exec::batch(rx, cmd, show_filesystem_errors)
            } else {
                let shared_rx = Arc::new(Mutex::new(rx));

                let out_perm = Arc::new(Mutex::new(()));

                // Each spawned job will store it's thread handle in here.
                let mut handles = Vec::with_capacity(threads);
                for _ in 0..threads {
                    let rx = Arc::clone(&shared_rx);
                    let cmd = Arc::clone(cmd);
                    let out_perm = Arc::clone(&out_perm);

                    // Spawn a job thread that will listen for and execute inputs.
                    let handle =
                        thread::spawn(move || exec::job(rx, cmd, out_perm, show_filesystem_errors));

                    // Push the handle of the spawned thread into the vector for later joining.
                    handles.push(handle);
                }

                // Wait for all threads to exit before exiting the program.
                for h in handles {
                    h.join().unwrap();
                }

                ExitCode::Success
            }
        } else {
            let start = time::Instant::now();

            let mut buffer = vec![];

            // Start in buffering mode
            let mut mode = ReceiverMode::Buffering;

            // Maximum time to wait before we start streaming to the console.
            let max_buffer_time = config
                .max_buffer_time
                .unwrap_or_else(|| time::Duration::from_millis(100));

            let stdout = io::stdout();
            let mut stdout = stdout.lock();

            for worker_result in rx {
                match worker_result {
                    WorkerResult::Entry(value) => {
                        match mode {
                            ReceiverMode::Buffering => {
                                buffer.push(value);

                                // Have we reached the maximum buffer size or maximum buffering time?
                                if buffer.len() > MAX_BUFFER_LENGTH
                                    || time::Instant::now() - start > max_buffer_time
                                {
                                    // Flush the buffer
                                    for v in &buffer {
                                        output::print_entry(
                                            &mut stdout,
                                            v,
                                            &config,
                                            &wants_to_quit,
                                        );
                                    }
                                    buffer.clear();

                                    // Start streaming
                                    mode = ReceiverMode::Streaming;
                                }
                            }
                            ReceiverMode::Streaming => {
                                output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                            }
                        }
                    }
                    WorkerResult::Error(err) => {
                        if show_filesystem_errors {
                            print_error!("{}", err);
                        }
                    }
                }
            }

            // If we have finished fast enough (faster than max_buffer_time), we haven't streamed
            // anything to the console, yet. In this case, sort the results and print them:
            if !buffer.is_empty() {
                buffer.sort();
                for value in buffer {
                    output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                }
            }

            ExitCode::Success
        }
    })
}

fn spawn_senders(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    pattern: Arc<Regex>,
    parallel_walker: ignore::WalkParallel,
    tx: Sender<WorkerResult>,
) {
    parallel_walker.run(|| {
        let config = Arc::clone(config);
        let pattern = Arc::clone(&pattern);
        let tx_thread = tx.clone();
        let wants_to_quit = Arc::clone(wants_to_quit);

        Box::new(move |entry_o| {
            if wants_to_quit.load(Ordering::Relaxed) {
                return ignore::WalkState::Quit;
            }

            let entry = match entry_o {
                Ok(e) => e,
                Err(err) => {
                    tx_thread.send(WorkerResult::Error(err)).unwrap();
                    return ignore::WalkState::Continue;
                }
            };

            if entry.depth() == 0 {
                return ignore::WalkState::Continue;
            }

            // Check the name first, since it doesn't require metadata
            let entry_path = entry.path();

            let search_str: Cow<OsStr> = if config.search_full_path {
                match fshelper::path_absolute_form(entry_path) {
                    Ok(path_abs_buf) => Cow::Owned(path_abs_buf.as_os_str().to_os_string()),
                    Err(_) => {
                        print_error_and_exit!("Unable to retrieve absolute path.");
                    }
                }
            } else {
                match entry_path.file_name() {
                    Some(filename) => Cow::Borrowed(filename),
                    None => unreachable!(
                        "Encountered file system entry without a file name. This should only \
                         happen for paths like 'foo/bar/..' or '/' which are not supposed to \
                         appear in a file system traversal."
                    ),
                }
            };

            if !pattern.is_match(&osstr_to_bytes(search_str.as_ref())) {
                return ignore::WalkState::Continue;
            }

            // Filter out unwanted extensions.
            if let Some(ref exts_regex) = config.extensions {
                if let Some(path_str) = entry_path.file_name() {
                    if !exts_regex.is_match(&osstr_to_bytes(path_str)) {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted file types.
            if let Some(ref file_types) = config.file_types {
                if let Some(ref entry_type) = entry.file_type() {
                    if (!file_types.files && entry_type.is_file())
                        || (!file_types.directories && entry_type.is_dir())
                        || (!file_types.symlinks && entry_type.is_symlink())
                        || (file_types.executables_only
                            && !entry
                                .metadata()
                                .map(|m| fshelper::is_executable(&m))
                                .unwrap_or(false))
                        || (file_types.empty_only && !fshelper::is_empty(&entry))
                    {
                        return ignore::WalkState::Continue;
                    } else if !(entry_type.is_file()
                        || entry_type.is_dir()
                        || entry_type.is_symlink())
                    {
                        // This is probably a block device, char device, fifo or socket. Skip it.
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted sizes if it is a file and we have been given size constraints.
            if !config.size_constraints.is_empty() {
                if entry_path.is_file() {
                    if let Ok(metadata) = entry_path.metadata() {
                        let file_size = metadata.len();
                        if config
                            .size_constraints
                            .iter()
                            .any(|sc| !sc.is_within(file_size))
                        {
                            return ignore::WalkState::Continue;
                        }
                    } else {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted modification times
            if !config.time_constraints.is_empty() {
                let mut matched = false;
                if entry_path.is_file() {
                    if let Ok(metadata) = entry_path.metadata() {
                        if let Ok(modified) = metadata.modified() {
                            matched = config
                                .time_constraints
                                .iter()
                                .all(|tf| tf.applies_to(&modified));
                        }
                    }
                }
                if !matched {
                    return ignore::WalkState::Continue;
                }
            }

            // TODO: take care of the unwrap call
            tx_thread
                .send(WorkerResult::Entry(entry_path.to_owned()))
                .unwrap();

            ignore::WalkState::Continue
        })
    });
}
<commit_msg>Make --changed-within/before work for directories<commit_after>// Copyright (c) 2017 fd developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use crate::exec;
use crate::exit_codes::ExitCode;
use crate::fshelper;
use crate::internal::{opts::FdOptions, osstr_to_bytes, MAX_BUFFER_LENGTH};
use crate::output;

use std::borrow::Cow;
use std::error::Error;
use std::ffi::OsStr;
use std::io;
use std::path::PathBuf;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use ignore::overrides::OverrideBuilder;
use ignore::{self, WalkBuilder};
use regex::bytes::Regex;

/// The receiver thread can either be buffering results or directly streaming to the console.
enum ReceiverMode {
    /// Receiver is still buffering in order to sort the results, if the search finishes fast
    /// enough.
    Buffering,

    /// Receiver is directly printing results to the output.
    Streaming,
}

/// The Worker threads can result in a valid entry having PathBuf or an error.
pub enum WorkerResult {
    Entry(PathBuf),
    Error(ignore::Error),
}

/// Recursively scan the given search path for files / pathnames matching the pattern.
///
/// If the `--exec` argument was supplied, this will create a thread pool for executing
/// jobs in parallel from a given command line and the discovered paths. Otherwise, each
/// path will simply be written to standard output.
pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<FdOptions>) -> ExitCode {
    let mut path_iter = path_vec.iter();
    let first_path_buf = path_iter
        .next()
        .expect("Error: Path vector can not be empty");
    let (tx, rx) = channel();

    let mut override_builder = OverrideBuilder::new(first_path_buf.as_path());

    for pattern in &config.exclude_patterns {
        let res = override_builder.add(pattern);
        if res.is_err() {
            print_error_and_exit!("Malformed exclude pattern '{}'", pattern);
        }
    }
    let overrides = override_builder.build().unwrap_or_else(|_| {
        print_error_and_exit!("Mismatch in exclude patterns");
    });

    let mut walker = WalkBuilder::new(first_path_buf.as_path());
    walker
        .hidden(config.ignore_hidden)
        .ignore(config.read_fdignore)
        .parents(config.read_fdignore || config.read_vcsignore)
        .git_ignore(config.read_vcsignore)
        .git_global(config.read_vcsignore)
        .git_exclude(config.read_vcsignore)
        .overrides(overrides)
        .follow_links(config.follow_links)
        .max_depth(config.max_depth);

    if config.read_fdignore {
        walker.add_custom_ignore_filename(".fdignore");
    }

    for ignore_file in &config.ignore_files {
        let result = walker.add_ignore(ignore_file);
        match result {
            Some(ignore::Error::Partial(_)) => (),
            Some(err) => {
                print_error!(
                    "{}",
                    format!(
                        "Malformed pattern in custom ignore file '{}': {}.",
                        ignore_file.to_string_lossy(),
                        err.description()
                    )
                );
            }
            None => (),
        }
    }

    for path_entry in path_iter {
        walker.add(path_entry.as_path());
    }

    let parallel_walker = walker.threads(config.threads).build_parallel();

    let wants_to_quit = Arc::new(AtomicBool::new(false));
    if config.ls_colors.is_some() && config.command.is_none() {
        let wq = Arc::clone(&wants_to_quit);
        ctrlc::set_handler(move || {
            if wq.load(Ordering::Relaxed) {
                // Ctrl-C has been pressed twice, exit NOW
                process::exit(ExitCode::KilledBySigint.into());
            } else {
                wq.store(true, Ordering::Relaxed);
            }
        })
        .unwrap();
    }

    // Spawn the thread that receives all results through the channel.
    let receiver_thread = spawn_receiver(&config, &wants_to_quit, rx);

    // Spawn the sender threads.
    spawn_senders(&config, &wants_to_quit, pattern, parallel_walker, tx);

    // Wait for the receiver thread to print out all results.
    let exit_code = receiver_thread.join().unwrap();

    if wants_to_quit.load(Ordering::Relaxed) {
        process::exit(ExitCode::KilledBySigint.into());
    }

    exit_code
}

fn spawn_receiver(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    rx: Receiver<WorkerResult>,
) -> thread::JoinHandle<ExitCode> {
    let config = Arc::clone(config);
    let wants_to_quit = Arc::clone(wants_to_quit);

    let show_filesystem_errors = config.show_filesystem_errors;
    let threads = config.threads;

    thread::spawn(move || {
        // This will be set to `Some` if the `--exec` argument was supplied.
        if let Some(ref cmd) = config.command {
            if cmd.in_batch_mode() {
                exec::batch(rx, cmd, show_filesystem_errors)
            } else {
                let shared_rx = Arc::new(Mutex::new(rx));

                let out_perm = Arc::new(Mutex::new(()));

                // Each spawned job will store it's thread handle in here.
                let mut handles = Vec::with_capacity(threads);
                for _ in 0..threads {
                    let rx = Arc::clone(&shared_rx);
                    let cmd = Arc::clone(cmd);
                    let out_perm = Arc::clone(&out_perm);

                    // Spawn a job thread that will listen for and execute inputs.
                    let handle =
                        thread::spawn(move || exec::job(rx, cmd, out_perm, show_filesystem_errors));

                    // Push the handle of the spawned thread into the vector for later joining.
                    handles.push(handle);
                }

                // Wait for all threads to exit before exiting the program.
                for h in handles {
                    h.join().unwrap();
                }

                ExitCode::Success
            }
        } else {
            let start = time::Instant::now();

            let mut buffer = vec![];

            // Start in buffering mode
            let mut mode = ReceiverMode::Buffering;

            // Maximum time to wait before we start streaming to the console.
            let max_buffer_time = config
                .max_buffer_time
                .unwrap_or_else(|| time::Duration::from_millis(100));

            let stdout = io::stdout();
            let mut stdout = stdout.lock();

            for worker_result in rx {
                match worker_result {
                    WorkerResult::Entry(value) => {
                        match mode {
                            ReceiverMode::Buffering => {
                                buffer.push(value);

                                // Have we reached the maximum buffer size or maximum buffering time?
                                if buffer.len() > MAX_BUFFER_LENGTH
                                    || time::Instant::now() - start > max_buffer_time
                                {
                                    // Flush the buffer
                                    for v in &buffer {
                                        output::print_entry(
                                            &mut stdout,
                                            v,
                                            &config,
                                            &wants_to_quit,
                                        );
                                    }
                                    buffer.clear();

                                    // Start streaming
                                    mode = ReceiverMode::Streaming;
                                }
                            }
                            ReceiverMode::Streaming => {
                                output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                            }
                        }
                    }
                    WorkerResult::Error(err) => {
                        if show_filesystem_errors {
                            print_error!("{}", err);
                        }
                    }
                }
            }

            // If we have finished fast enough (faster than max_buffer_time), we haven't streamed
            // anything to the console, yet. In this case, sort the results and print them:
            if !buffer.is_empty() {
                buffer.sort();
                for value in buffer {
                    output::print_entry(&mut stdout, &value, &config, &wants_to_quit);
                }
            }

            ExitCode::Success
        }
    })
}

fn spawn_senders(
    config: &Arc<FdOptions>,
    wants_to_quit: &Arc<AtomicBool>,
    pattern: Arc<Regex>,
    parallel_walker: ignore::WalkParallel,
    tx: Sender<WorkerResult>,
) {
    parallel_walker.run(|| {
        let config = Arc::clone(config);
        let pattern = Arc::clone(&pattern);
        let tx_thread = tx.clone();
        let wants_to_quit = Arc::clone(wants_to_quit);

        Box::new(move |entry_o| {
            if wants_to_quit.load(Ordering::Relaxed) {
                return ignore::WalkState::Quit;
            }

            let entry = match entry_o {
                Ok(e) => e,
                Err(err) => {
                    tx_thread.send(WorkerResult::Error(err)).unwrap();
                    return ignore::WalkState::Continue;
                }
            };

            if entry.depth() == 0 {
                return ignore::WalkState::Continue;
            }

            // Check the name first, since it doesn't require metadata
            let entry_path = entry.path();

            let search_str: Cow<OsStr> = if config.search_full_path {
                match fshelper::path_absolute_form(entry_path) {
                    Ok(path_abs_buf) => Cow::Owned(path_abs_buf.as_os_str().to_os_string()),
                    Err(_) => {
                        print_error_and_exit!("Unable to retrieve absolute path.");
                    }
                }
            } else {
                match entry_path.file_name() {
                    Some(filename) => Cow::Borrowed(filename),
                    None => unreachable!(
                        "Encountered file system entry without a file name. This should only \
                         happen for paths like 'foo/bar/..' or '/' which are not supposed to \
                         appear in a file system traversal."
                    ),
                }
            };

            if !pattern.is_match(&osstr_to_bytes(search_str.as_ref())) {
                return ignore::WalkState::Continue;
            }

            // Filter out unwanted extensions.
            if let Some(ref exts_regex) = config.extensions {
                if let Some(path_str) = entry_path.file_name() {
                    if !exts_regex.is_match(&osstr_to_bytes(path_str)) {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted file types.
            if let Some(ref file_types) = config.file_types {
                if let Some(ref entry_type) = entry.file_type() {
                    if (!file_types.files && entry_type.is_file())
                        || (!file_types.directories && entry_type.is_dir())
                        || (!file_types.symlinks && entry_type.is_symlink())
                        || (file_types.executables_only
                            && !entry
                                .metadata()
                                .map(|m| fshelper::is_executable(&m))
                                .unwrap_or(false))
                        || (file_types.empty_only && !fshelper::is_empty(&entry))
                    {
                        return ignore::WalkState::Continue;
                    } else if !(entry_type.is_file()
                        || entry_type.is_dir()
                        || entry_type.is_symlink())
                    {
                        // This is probably a block device, char device, fifo or socket. Skip it.
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted sizes if it is a file and we have been given size constraints.
            if !config.size_constraints.is_empty() {
                if entry_path.is_file() {
                    if let Ok(metadata) = entry_path.metadata() {
                        let file_size = metadata.len();
                        if config
                            .size_constraints
                            .iter()
                            .any(|sc| !sc.is_within(file_size))
                        {
                            return ignore::WalkState::Continue;
                        }
                    } else {
                        return ignore::WalkState::Continue;
                    }
                } else {
                    return ignore::WalkState::Continue;
                }
            }

            // Filter out unwanted modification times
            if !config.time_constraints.is_empty() {
                let mut matched = false;
                if let Ok(metadata) = entry_path.metadata() {
                    if let Ok(modified) = metadata.modified() {
                        matched = config
                            .time_constraints
                            .iter()
                            .all(|tf| tf.applies_to(&modified));
                    }
                }
                if !matched {
                    return ignore::WalkState::Continue;
                }
            }

            // TODO: take care of the unwrap call
            tx_thread
                .send(WorkerResult::Entry(entry_path.to_owned()))
                .unwrap();

            ignore::WalkState::Continue
        })
    });
}
<|endoftext|>