<commit_before>mod reader;
mod writer;

use std::time::{Duration, Instant};
use std::{io, mem};
use std::net::SocketAddr;
use std::sync::Arc;

use url::percent_encoding::percent_encode_byte;
use url::Url;

use {amy, bencode, PEER_ID};
use self::writer::Writer;
use self::reader::{ReadRes, Reader};
use socket::TSocket;
use tracker::{self, dns, Announce, Error, ErrorKind, Response, Result, ResultExt, TrackerResponse};
use util::UHashMap;

const TIMEOUT_MS: u64 = 5_000;

pub struct Handler {
    reg: amy::Registrar,
    connections: UHashMap<Tracker>,
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    url: Arc<Url>,
    last_updated: Instant,
    redirect: bool,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS {
        sock: TSocket,
        req: Vec<u8>,
        port: u16,
    },
    Writing {
        sock: TSocket,
        writer: Writer,
    },
    Reading {
        sock: TSocket,
        reader: Reader,
    },
    Redirect(String),
    Complete(TrackerResponse),
}

enum HTTPRes {
    None,
    Redirect(String),
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<HTTPRes> {
        let s = mem::replace(self, TrackerState::Error);
        match s.next(event)? {
            TrackerState::Complete(r) => Ok(HTTPRes::Complete(r)),
            TrackerState::Redirect(l) => Ok(HTTPRes::Redirect(l)),
            n => {
                mem::replace(self, n);
                Ok(HTTPRes::None)
            }
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (
                TrackerState::ResolvingDNS {
                    mut sock,
                    req,
                    port,
                },
                Event::DNSResolved(r),
            ) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr).chain_err(|| ErrorKind::IO)?;
                Ok(TrackerState::Writing {
                    sock,
                    writer: Writer::new(req),
                }.next(Event::Writable)?
                    .next(Event::Readable)?)
            }
            (
                TrackerState::Writing {
                    mut sock,
                    mut writer,
                },
                _,
            ) => match writer.writable(&mut sock)? {
                Some(()) => {
                    debug!("Tracker write completed, beginning read");
                    let r = Reader::new();
                    Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                }
                None => Ok(TrackerState::Writing { sock, writer }),
            },
            (
                TrackerState::Reading {
                    mut sock,
                    mut reader,
                },
                _,
            ) => match reader.readable(&mut sock)? {
                ReadRes::Done(data) => {
                    let content = bencode::decode_buf(&data)
                        .chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                }
                ReadRes::Redirect(l) => Ok(TrackerState::Redirect(l)),
                ReadRes::None => Ok(TrackerState::Reading { sock, reader }),
            },
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!"),
        }
    }
}

impl Handler {
    pub fn new(reg: &amy::Registrar) -> io::Result<Handler> {
        Ok(Handler {
            reg: reg.clone(),
            connections: UHashMap::default(),
        })
    }

    pub fn active_requests(&self) -> usize {
        self.connections.len()
    }

    pub fn complete(&self) -> bool {
        self.connections.is_empty()
    }

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        let id = resp.id;
        debug!("Received a DNS resp for {:?}", id);
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn readable(&mut self, id: usize, dns: &mut dns::Resolver) -> Option<Response> {
        let mut loc = None;
        let mut resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(HTTPRes::Complete(r)) => {
                    debug!("Announce response received for {:?} succesfully", id);
                    Some(Response::Tracker {
                        tid: trk.torrent,
                        url: trk.url.clone(),
                        resp: Ok(r),
                    })
                }
                Ok(HTTPRes::Redirect(l)) => {
                    loc = Some((l, trk.url.clone()));
                    None
                }
                Ok(HTTPRes::None) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };

        if resp.is_some() {
            self.connections.remove(&id);
        }

        if let Some((l, old)) = loc {
            let trk = self.connections.remove(&id).unwrap();
            // Disallow 2 levels of redirection
            if trk.redirect {
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::InvalidResponse("Too many redirects").into()),
                });
            }
            if let Err(e) = self.try_redirect(&l, old, trk.torrent, dns) {
                debug!(
                    "Announce response received for {:?}, redirecting!",
                    trk.torrent
                );
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                });
            }
        }
        resp
    }

    fn try_redirect(
        &mut self,
        url: &str,
        original_url: Arc<Url>,
        torrent: usize,
        dns: &mut dns::Resolver,
    ) -> Result<()> {
        let url = Url::parse(url).chain_err(|| ErrorKind::InvalidResponse("Malformed redirect!"))?;
        let mut http_req = Vec::with_capacity(50);
        http_req.extend_from_slice(b"GET ");
        http_req.extend_from_slice(url.path().as_bytes());
        if let Some(q) = url.query() {
            http_req.extend_from_slice(b"?");
            http_req.extend_from_slice(q.as_bytes());
        }

        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        let user_agent = format!(
            "User-Agent: {}/{}\r\n",
            "synapse",
            env!("CARGO_PKG_VERSION")
        );
        http_req.extend_from_slice(user_agent.as_bytes());
        http_req.extend_from_slice(b"Connection: close\r\n");
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str()
            .ok_or_else(|| Error::from(ErrorKind::InvalidResponse("Malformed redirect!")))?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n\r\n");

        let ohost = if url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                last_updated: Instant::now(),
                redirect: true,
                torrent,
                url: original_url,
                state: TrackerState::new(sock, http_req, port),
            },
        );

        debug!("Dispatching redirect DNS req, id {:?}", id);
        if let Some(ip) = dns.new_query(id, host).chain_err(|| ErrorKind::IO)? {
            debug!("Using cached DNS response");
            let res = self.dns_resolved(dns::QueryResponse { id, res: Ok(ip) });
            if res.is_some() {
                bail!("Failed to establish connection to tracker!");
            }
        }
        Ok(())
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        self.connections.retain(|id, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                debug!("Announce {:?} timed out", id);
                resps.push(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::Timeout.into()),
                });
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn new_announce(&mut self, req: Announce, dns: &mut dns::Resolver) -> Result<()> {
        debug!("Received a new announce req for {:?}", req.url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL
        http_req.extend_from_slice(req.url.path().as_bytes());
        http_req.extend_from_slice(b"?");
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        if let Some(nw) = req.num_want {
            append_query_pair(&mut http_req, "numwant", &nw.to_string());
        }
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "stopped");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {}
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        let user_agent = format!(
            "User-Agent: {}/{}\r\n",
            "synapse",
            env!("CARGO_PKG_VERSION")
        );
        http_req.extend_from_slice(user_agent.as_bytes());
        // Don't keep alive
        http_req.extend_from_slice(b"Connection: close\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = req.url.host_str().ok_or_else(|| {
            Error::from(ErrorKind::InvalidRequest(
                "Tracker announce url has no host!".to_owned(),
            ))
        })?;
        let port =
            req.url
                .port()
                .unwrap_or_else(|| if req.url.scheme() == "https" { 443 } else { 80 });
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        let ohost = if req.url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                url: req.url.clone(),
                last_updated: Instant::now(),
                torrent: req.id,
                state: TrackerState::new(sock, http_req, port),
                redirect: false,
            },
        );

        debug!("Dispatching DNS req, id {:?}", id);
        if let Some(ip) = dns.new_query(id, host).chain_err(|| ErrorKind::IO)? {
            debug!("Using cached DNS response");
            let res = self.dns_resolved(dns::QueryResponse { id, res: Ok(ip) });
            if res.is_some() {
                bail!("Failed to estbalish connection to tracker!");
            }
        }

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice(b"=");
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice(b"&");
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        let c = char::from(*byte);
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            resp.push(c);
        } else {
            resp += percent_encode_byte(*byte);
        }
    }
    resp
}
<commit_msg>Add in tracker query pairs to announce<commit_after>mod reader;
mod writer;

use std::time::{Duration, Instant};
use std::{io, mem};
use std::net::SocketAddr;
use std::sync::Arc;

use url::percent_encoding::percent_encode_byte;
use url::Url;

use {amy, bencode, PEER_ID};
use self::writer::Writer;
use self::reader::{ReadRes, Reader};
use socket::TSocket;
use tracker::{self, dns, Announce, Error, ErrorKind, Response, Result, ResultExt, TrackerResponse};
use util::UHashMap;

const TIMEOUT_MS: u64 = 5_000;

pub struct Handler {
    reg: amy::Registrar,
    connections: UHashMap<Tracker>,
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    url: Arc<Url>,
    last_updated: Instant,
    redirect: bool,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS {
        sock: TSocket,
        req: Vec<u8>,
        port: u16,
    },
    Writing {
        sock: TSocket,
        writer: Writer,
    },
    Reading {
        sock: TSocket,
        reader: Reader,
    },
    Redirect(String),
    Complete(TrackerResponse),
}

enum HTTPRes {
    None,
    Redirect(String),
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<HTTPRes> {
        let s = mem::replace(self, TrackerState::Error);
        match s.next(event)? {
            TrackerState::Complete(r) => Ok(HTTPRes::Complete(r)),
            TrackerState::Redirect(l) => Ok(HTTPRes::Redirect(l)),
            n => {
                mem::replace(self, n);
                Ok(HTTPRes::None)
            }
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (
                TrackerState::ResolvingDNS {
                    mut sock,
                    req,
                    port,
                },
                Event::DNSResolved(r),
            ) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr).chain_err(|| ErrorKind::IO)?;
                Ok(TrackerState::Writing {
                    sock,
                    writer: Writer::new(req),
                }.next(Event::Writable)?
                    .next(Event::Readable)?)
            }
            (
                TrackerState::Writing {
                    mut sock,
                    mut writer,
                },
                _,
            ) => match writer.writable(&mut sock)? {
                Some(()) => {
                    debug!("Tracker write completed, beginning read");
                    let r = Reader::new();
                    Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                }
                None => Ok(TrackerState::Writing { sock, writer }),
            },
            (
                TrackerState::Reading {
                    mut sock,
                    mut reader,
                },
                _,
            ) => match reader.readable(&mut sock)? {
                ReadRes::Done(data) => {
                    let content = bencode::decode_buf(&data)
                        .chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                }
                ReadRes::Redirect(l) => Ok(TrackerState::Redirect(l)),
                ReadRes::None => Ok(TrackerState::Reading { sock, reader }),
            },
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!"),
        }
    }
}

impl Handler {
    pub fn new(reg: &amy::Registrar) -> io::Result<Handler> {
        Ok(Handler {
            reg: reg.clone(),
            connections: UHashMap::default(),
        })
    }

    pub fn active_requests(&self) -> usize {
        self.connections.len()
    }

    pub fn complete(&self) -> bool {
        self.connections.is_empty()
    }

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        let id = resp.id;
        debug!("Received a DNS resp for {:?}", id);
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        let resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };
        if resp.is_some() {
            self.connections.remove(&id);
        }
        resp
    }

    pub fn readable(&mut self, id: usize, dns: &mut dns::Resolver) -> Option<Response> {
        let mut loc = None;
        let mut resp = if let Some(trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(HTTPRes::Complete(r)) => {
                    debug!("Announce response received for {:?} succesfully", id);
                    Some(Response::Tracker {
                        tid: trk.torrent,
                        url: trk.url.clone(),
                        resp: Ok(r),
                    })
                }
                Ok(HTTPRes::Redirect(l)) => {
                    loc = Some((l, trk.url.clone()));
                    None
                }
                Ok(HTTPRes::None) => None,
                Err(e) => Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                }),
            }
        } else {
            None
        };

        if resp.is_some() {
            self.connections.remove(&id);
        }

        if let Some((l, old)) = loc {
            let trk = self.connections.remove(&id).unwrap();
            // Disallow 2 levels of redirection
            if trk.redirect {
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::InvalidResponse("Too many redirects").into()),
                });
            }
            if let Err(e) = self.try_redirect(&l, old, trk.torrent, dns) {
                debug!(
                    "Announce response received for {:?}, redirecting!",
                    trk.torrent
                );
                resp = Some(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(e),
                });
            }
        }
        resp
    }

    fn try_redirect(
        &mut self,
        url: &str,
        original_url: Arc<Url>,
        torrent: usize,
        dns: &mut dns::Resolver,
    ) -> Result<()> {
        let url = Url::parse(url).chain_err(|| ErrorKind::InvalidResponse("Malformed redirect!"))?;
        let mut http_req = Vec::with_capacity(50);
        http_req.extend_from_slice(b"GET ");
        http_req.extend_from_slice(url.path().as_bytes());
        if let Some(q) = url.query() {
            http_req.extend_from_slice(b"?");
            http_req.extend_from_slice(q.as_bytes());
        }

        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        let user_agent = format!(
            "User-Agent: {}/{}\r\n",
            "synapse",
            env!("CARGO_PKG_VERSION")
        );
        http_req.extend_from_slice(user_agent.as_bytes());
        http_req.extend_from_slice(b"Connection: close\r\n");
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str()
            .ok_or_else(|| Error::from(ErrorKind::InvalidResponse("Malformed redirect!")))?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n\r\n");

        let ohost = if url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                last_updated: Instant::now(),
                redirect: true,
                torrent,
                url: original_url,
                state: TrackerState::new(sock, http_req, port),
            },
        );

        debug!("Dispatching redirect DNS req, id {:?}", id);
        if let Some(ip) = dns.new_query(id, host).chain_err(|| ErrorKind::IO)? {
            debug!("Using cached DNS response");
            let res = self.dns_resolved(dns::QueryResponse { id, res: Ok(ip) });
            if res.is_some() {
                bail!("Failed to establish connection to tracker!");
            }
        }
        Ok(())
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        self.connections.retain(|id, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                debug!("Announce {:?} timed out", id);
                resps.push(Response::Tracker {
                    tid: trk.torrent,
                    url: trk.url.clone(),
                    resp: Err(ErrorKind::Timeout.into()),
                });
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn new_announce(&mut self, req: Announce, dns: &mut dns::Resolver) -> Result<()> {
        debug!("Received a new announce req for {:?}", req.url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL
        http_req.extend_from_slice(req.url.path().as_bytes());
        http_req.extend_from_slice(b"?");
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        if let Some(nw) = req.num_want {
            append_query_pair(&mut http_req, "numwant", &nw.to_string());
        }
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "stopped");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {}
        }
        for (k, v) in req.url.query_pairs() {
            append_query_pair(&mut http_req, &k, &v);
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        let user_agent = format!(
            "User-Agent: {}/{}\r\n",
            "synapse",
            env!("CARGO_PKG_VERSION")
        );
        http_req.extend_from_slice(user_agent.as_bytes());
        // Don't keep alive
        http_req.extend_from_slice(b"Connection: close\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = req.url.host_str().ok_or_else(|| {
            Error::from(ErrorKind::InvalidRequest(
                "Tracker announce url has no host!".to_owned(),
            ))
        })?;
        let port =
            req.url
                .port()
                .unwrap_or_else(|| if req.url.scheme() == "https" { 443 } else { 80 });
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        let ohost = if req.url.scheme() == "https" {
            Some(host.to_owned())
        } else {
            None
        };

        // Setup actual connection and start DNS query
        let sock = TSocket::new_v4(ohost).chain_err(|| ErrorKind::IO)?;
        let id = self.reg
            .register(&sock, amy::Event::Both)
            .chain_err(|| ErrorKind::IO)?;
        self.connections.insert(
            id,
            Tracker {
                url: req.url.clone(),
                last_updated: Instant::now(),
                torrent: req.id,
                state: TrackerState::new(sock, http_req, port),
                redirect: false,
            },
        );

        debug!("Dispatching DNS req, id {:?}", id);
        if let Some(ip) = dns.new_query(id, host).chain_err(|| ErrorKind::IO)? {
            debug!("Using cached DNS response");
            let res = self.dns_resolved(dns::QueryResponse { id, res: Ok(ip) });
            if res.is_some() {
                bail!("Failed to estbalish connection to tracker!");
            }
        }

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice(b"=");
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice(b"&");
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        let c = char::from(*byte);
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            resp.push(c);
        } else {
            resp += percent_encode_byte(*byte);
        }
    }
    resp
}
<|endoftext|>


<commit_before>use std::hash::Hash;
use std::collections::{BTreeMap, HashSet};
use std::cmp::Ordering;
use parking_lot::RwLock;
use bifrost_hasher::hash_str;

#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
pub enum Relation {
    Equal,
    Before,
    After,
    Concurrent
}

#[derive(Serialize, Deserialize, Debug, Clone, Hash)]
pub struct VectorClock<S: Ord + Eq + Copy> {
    map: BTreeMap<S, u64>
}

impl <S: Eq + Copy + Ord> PartialOrd for VectorClock<S> {
    fn partial_cmp(&self, other: &VectorClock<S>) -> Option<Ordering> {
        let rel = self.relation(other);
        match rel {
            Relation::Before => Some(Ordering::Less),
            Relation::After => Some(Ordering::Greater),
            Relation::Equal => Some(Ordering::Equal),
            Relation::Concurrent => None
        }
    }

    fn lt(&self, other: &VectorClock<S>) -> bool {
        self.relation(other) == Relation::Before
    }

    fn le(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::Before || rel == Relation::Equal
    }

    fn gt(&self, other: &VectorClock<S>) -> bool {
        self.relation(other) == Relation::After
    }

    fn ge(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::After || rel == Relation::Equal
    }
}

impl <S: Eq + Copy + Ord> PartialEq for VectorClock<S> {
    fn eq(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::Equal
    }

    fn ne(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        !(rel == Relation::After || rel == Relation::Before)
    }
}

impl <S: Ord + Eq + Copy>VectorClock<S> {
    pub fn new() -> VectorClock<S> {
        VectorClock {
            map: BTreeMap::new()
        }
    }

    pub fn inc(&mut self, server: S) -> VectorClock<S> {
        *self.map.entry(server).or_insert(0) += 1;
        self.clone()
    }

    pub fn happened_before(&self, clock_b: &VectorClock<S>) -> bool {
        let mut a_lt_b = false;
        for (server, ai) in self.map.iter() {
            let bi = *clock_b.map.get(server).unwrap_or(&0);
            if *ai > bi {return false;}
            a_lt_b = a_lt_b || *ai < bi;
        }
        for (server, bi) in clock_b.map.iter() {
            let ai = *self.map.get(server).unwrap_or(&0);
            if ai > *bi {return false;}
            a_lt_b = a_lt_b || ai < *bi;
        }
        return a_lt_b;
    }
    pub fn equals(&self, clock_b: &VectorClock<S>) -> bool {
        for (server, ac) in self.map.iter() {
            if let Some(bc) = clock_b.map.get(server) {
                if ac != bc {return false;}
            }
        }
        return true;
    }
    pub fn relation(&self, clock_b: &VectorClock<S>) -> Relation {
        if self.equals(clock_b) {
            return Relation::Equal;
        }
        if self.happened_before(clock_b) {
            return Relation::Before;
        }
        if clock_b.happened_before(self) {
            return Relation::After;
        }
        return Relation::Concurrent;
    }
    pub fn merge_with(&mut self, clock_b: &VectorClock<S>) {
        // merge_with is used to update counter for other servers (also learn from it)
        for (server, bc) in clock_b.map.iter() {
            let mut ba = self.map.entry(*server).or_insert(0);
            if *ba < *bc {*ba = *bc}
        }
    }
    pub fn learn_from(&mut self, clock_b: &VectorClock<S>) {
        // learn_from only insert missing servers into the clock
        for (server, bc) in clock_b.map.iter() {
            self.map.entry(*server).or_insert(*bc);
        }
    }
}

pub struct ServerVectorClock {
    server: u64,
    clock: RwLock<VectorClock<u64>>
}

impl ServerVectorClock {
    pub fn new(server_address: &String) -> ServerVectorClock {
        ServerVectorClock {
            server: hash_str(server_address),
            clock: RwLock::new(VectorClock::new())
        }
    }
    pub fn inc(&self) -> StandardVectorClock {
        let mut clock = self.clock.write();
        clock.inc(self.server)
    }

    pub fn happened_before(&self, clock_b: &StandardVectorClock) -> bool {
        let clock = self.clock.read();
        clock.happened_before(clock_b)
    }
    pub fn equals(&self, clock_b: &StandardVectorClock) -> bool {
        let clock = self.clock.read();
        clock.equals(clock_b)
    }
    pub fn relation(&self, clock_b: &StandardVectorClock) -> Relation {
        let clock = self.clock.read();
        clock.relation(clock_b)
    }
    pub fn merge_with(&self, clock_b: &StandardVectorClock) {
        let mut clock = self.clock.write();
        clock.merge_with(clock_b)
    }
    pub fn learn_from(&self, clock_b: &StandardVectorClock) {
        let mut clock = self.clock.write();
        clock.learn_from(clock_b)
    }
    pub fn to_clock(&self) -> StandardVectorClock {
        let clock = self.clock.read();
        clock.clone()
    }
}

pub type StandardVectorClock = VectorClock<u64>;<commit_msg>impl Eq for VectorClock<commit_after>use std::hash::Hash;
use std::collections::{BTreeMap, HashSet};
use std::cmp::Ordering;
use parking_lot::RwLock;
use bifrost_hasher::hash_str;

#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
pub enum Relation {
    Equal,
    Before,
    After,
    Concurrent
}

#[derive(Serialize, Deserialize, Debug, Clone, Hash, Eq)]
pub struct VectorClock<S: Ord + Eq + Copy> {
    map: BTreeMap<S, u64>
}

impl <S: Eq + Copy + Ord> PartialOrd for VectorClock<S> {
    fn partial_cmp(&self, other: &VectorClock<S>) -> Option<Ordering> {
        let rel = self.relation(other);
        match rel {
            Relation::Before => Some(Ordering::Less),
            Relation::After => Some(Ordering::Greater),
            Relation::Equal => Some(Ordering::Equal),
            Relation::Concurrent => None
        }
    }

    fn lt(&self, other: &VectorClock<S>) -> bool {
        self.relation(other) == Relation::Before
    }

    fn le(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::Before || rel == Relation::Equal
    }

    fn gt(&self, other: &VectorClock<S>) -> bool {
        self.relation(other) == Relation::After
    }

    fn ge(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::After || rel == Relation::Equal
    }
}

impl <S: Eq + Copy + Ord> PartialEq for VectorClock<S> {
    fn eq(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        rel == Relation::Equal
    }

    fn ne(&self, other: &VectorClock<S>) -> bool {
        let rel = self.relation(other);
        !(rel == Relation::After || rel == Relation::Before)
    }
}

impl <S: Ord + Eq + Copy>VectorClock<S> {
    pub fn new() -> VectorClock<S> {
        VectorClock {
            map: BTreeMap::new()
        }
    }

    pub fn inc(&mut self, server: S) -> VectorClock<S> {
        *self.map.entry(server).or_insert(0) += 1;
        self.clone()
    }

    pub fn happened_before(&self, clock_b: &VectorClock<S>) -> bool {
        let mut a_lt_b = false;
        for (server, ai) in self.map.iter() {
            let bi = *clock_b.map.get(server).unwrap_or(&0);
            if *ai > bi {return false;}
            a_lt_b = a_lt_b || *ai < bi;
        }
        for (server, bi) in clock_b.map.iter() {
            let ai = *self.map.get(server).unwrap_or(&0);
            if ai > *bi {return false;}
            a_lt_b = a_lt_b || ai < *bi;
        }
        return a_lt_b;
    }
    pub fn equals(&self, clock_b: &VectorClock<S>) -> bool {
        for (server, ac) in self.map.iter() {
            if let Some(bc) = clock_b.map.get(server) {
                if ac != bc {return false;}
            }
        }
        return true;
    }
    pub fn relation(&self, clock_b: &VectorClock<S>) -> Relation {
        if self.equals(clock_b) {
            return Relation::Equal;
        }
        if self.happened_before(clock_b) {
            return Relation::Before;
        }
        if clock_b.happened_before(self) {
            return Relation::After;
        }
        return Relation::Concurrent;
    }
    pub fn merge_with(&mut self, clock_b: &VectorClock<S>) {
        // merge_with is used to update counter for other servers (also learn from it)
        for (server, bc) in clock_b.map.iter() {
            let mut ba = self.map.entry(*server).or_insert(0);
            if *ba < *bc {*ba = *bc}
        }
    }
    pub fn learn_from(&mut self, clock_b: &VectorClock<S>) {
        // learn_from only insert missing servers into the clock
        for (server, bc) in clock_b.map.iter() {
            self.map.entry(*server).or_insert(*bc);
        }
    }
}

pub struct ServerVectorClock {
    server: u64,
    clock: RwLock<VectorClock<u64>>
}

impl ServerVectorClock {
    pub fn new(server_address: &String) -> ServerVectorClock {
        ServerVectorClock {
            server: hash_str(server_address),
            clock: RwLock::new(VectorClock::new())
        }
    }
    pub fn inc(&self) -> StandardVectorClock {
        let mut clock = self.clock.write();
        clock.inc(self.server)
    }

    pub fn happened_before(&self, clock_b: &StandardVectorClock) -> bool {
        let clock = self.clock.read();
        clock.happened_before(clock_b)
    }
    pub fn equals(&self, clock_b: &StandardVectorClock) -> bool {
        let clock = self.clock.read();
        clock.equals(clock_b)
    }
    pub fn relation(&self, clock_b: &StandardVectorClock) -> Relation {
        let clock = self.clock.read();
        clock.relation(clock_b)
    }
    pub fn merge_with(&self, clock_b: &StandardVectorClock) {
        let mut clock = self.clock.write();
        clock.merge_with(clock_b)
    }
    pub fn learn_from(&self, clock_b: &StandardVectorClock) {
        let mut clock = self.clock.write();
        clock.learn_from(clock_b)
    }
    pub fn to_clock(&self) -> StandardVectorClock {
        let clock = self.clock.read();
        clock.clone()
    }
}

pub type StandardVectorClock = VectorClock<u64>;<|endoftext|>


<commit_before>use futures::Future;
use futures_cpupool::{Builder, CpuPool};
use num_cpus;

lazy_static! {
    static ref POOL: CpuPool = Builder::new()
        .name_prefix("Bifrost Shared Exec Pool".to_string())
        .pool_size(num_cpus::get())
        .create();
}

pub fn exec<F>(future: F) -> impl Future<Item = F::Item, Error = F::Error>
    where
        F: Future + Send + 'static,
        F::Item: Send + 'static,
        F::Error: Send + 'static,
{
    POOL.spawn(future)
}

pub fn wait<F>(future: F) -> Result<F::Item, F::Error>
    where
        F: Future + Send + 'static,
        F::Item: Send + 'static,
        F::Error: Send + 'static,
{
    exec(future).wait()
}<commit_msg>wrap with send<commit_after>use futures::prelude::*;
use futures_cpupool::{Builder, CpuPool};
use num_cpus;

lazy_static! {
    static ref POOL: CpuPool = Builder::new()
        .name_prefix("Bifrost Shared Exec Pool".to_string())
        .pool_size(num_cpus::get())
        .create();
}

struct SendFuture<F>
where
    F: Future + 'static,
    F::Item: Send + 'static,
    F::Error: Send + 'static,
{
    inner: F,
}

impl<F> Future for SendFuture<F>
where
    F: Future + 'static,
    F::Item: Send + 'static,
    F::Error: Send + 'static,
{
    type Item = F::Item;
    type Error = F::Error;

    fn poll(&mut self) -> Result<Async<<Self as Future>::Item>, <Self as Future>::Error> {
        self.inner.poll()
    }
}

unsafe impl<F> Send for SendFuture<F>
where
    F: Future + 'static,
    F::Item: Send + 'static,
    F::Error: Send + 'static,
{
}

pub fn exec<F>(future: F) -> impl Future<Item = F::Item, Error = F::Error>
where
    F: Future + 'static,
    F::Item: Send + 'static,
    F::Error: Send + 'static,
{
    POOL.spawn(SendFuture { inner: future })
}

pub fn wait<F>(future: F) -> Result<F::Item, F::Error>
where
    F: Future + 'static,
    F::Item: Send + 'static,
    F::Error: Send + 'static,
{
    exec(future).wait()
}
<|endoftext|>


<commit_before>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub fn walk<P: AsRef<Path>>(root: P) -> std::io::Result<Walk> {
    let metadata = std::fs::symlink_metadata(&root)?;

    let iter = std::fs::read_dir(&root)?;
    Ok(Walk {
        root: Some(WalkEntry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: vec!(ListDir::from_read_dir(iter)),
    })
}

pub struct Walk {
    // TODO: Add support for stopping at device boundaries.
    root: Option<WalkEntry>,
    pending: Vec<ListDir>,
}

pub struct WalkEntry {
    path: PathBuf,
    metadata: Metadata,
}

impl Walk {

    fn push(&mut self, entry: &WalkEntry) {
        match std::fs::read_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(ListDir::from_read_dir(iter));
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<WalkEntry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for Walk {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl ListDir {

    /// Converts a standard `ReadDir` iterator.
    fn from_read_dir(iter: std::fs::ReadDir) -> ListDir {
        ListDir {
            iter: iter,
        }
    }
}

impl std::iter::Iterator for ListDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(WalkEntry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    // TODO: Add test case for non-existing directory.

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let iter = std::fs::read_dir(&tempdir).unwrap();
        let mut iter = ListDir::from_read_dir(iter);

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let iter = std::fs::read_dir(&tempdir).unwrap();

        let mut results = ListDir::from_read_dir(iter).collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let iter = std::fs::read_dir(&tempdir).unwrap();

        let mut results = ListDir::from_read_dir(iter).collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let iter = std::fs::read_dir(&tempdir).unwrap();

        let mut results = ListDir::from_read_dir(iter).collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<commit_msg>Create an auxiliary function for listing dirs.<commit_after>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub fn walk<P: AsRef<Path>>(root: P) -> std::io::Result<Walk> {
    let metadata = std::fs::symlink_metadata(&root)?;

    Ok(Walk {
        root: Some(WalkEntry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: vec!(list_dir(&root)?),
    })
}

fn list_dir<P: AsRef<Path>>(path: P) -> std::io::Result<ListDir> {
    let iter = std::fs::read_dir(path)?;

    Ok(ListDir {
        iter: iter,
    })
}

pub struct Walk {
    // TODO: Add support for stopping at device boundaries.
    root: Option<WalkEntry>,
    pending: Vec<ListDir>,
}

pub struct WalkEntry {
    path: PathBuf,
    metadata: Metadata,
}

impl Walk {

    fn push(&mut self, entry: &WalkEntry) {
        match list_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(iter);
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<WalkEntry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for Walk {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl std::iter::Iterator for ListDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(WalkEntry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    // TODO: Add test case for non-existing directory.

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = list_dir(&tempdir).unwrap();

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<|endoftext|>


<commit_before>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub struct Entry {
    pub path: PathBuf,
    pub metadata: Metadata,
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> std::io::Result<WalkDir> {
    let metadata = std::fs::symlink_metadata(&root)?;
    let pending = vec!(list_dir(&root)?);

    Ok(WalkDir {
        root: Some(Entry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: pending,
    })
}

fn list_dir<P: AsRef<Path>>(path: P) -> std::io::Result<ListDir> {
    let iter = std::fs::read_dir(path)?;

    Ok(ListDir {
        iter: iter,
    })
}

pub struct WalkDir {
    // TODO: Add support for stopping at device boundaries.
    root: Option<Entry>,
    pending: Vec<ListDir>,
}

impl WalkDir {

    fn push(&mut self, entry: &Entry) {
        match list_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(iter);
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<Entry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for WalkDir {

    type Item = Entry;

    fn next(&mut self) -> Option<Entry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl std::iter::Iterator for ListDir {

    type Item = Entry;

    fn next(&mut self) -> Option<Entry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(Entry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    #[test]
    fn test_list_dir_non_existing() {
        let tempdir = tempfile::tempdir().unwrap();

        let iter = list_dir(tempdir.path().join("foo"));
        assert!(iter.is_err());
    }

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = list_dir(&tempdir).unwrap();

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_non_existing() {
        let tempdir = tempfile::tempdir().unwrap();

        let iter = walk_dir(tempdir.path().join("foo"));
        assert!(iter.is_err());
    }

    #[test]
    fn test_walk_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk_dir(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_dir_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_dir_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_dir_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_dir_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<commit_msg>Add directory listing test with unicode filenames.<commit_after>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub struct Entry {
    pub path: PathBuf,
    pub metadata: Metadata,
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> std::io::Result<WalkDir> {
    let metadata = std::fs::symlink_metadata(&root)?;
    let pending = vec!(list_dir(&root)?);

    Ok(WalkDir {
        root: Some(Entry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: pending,
    })
}

fn list_dir<P: AsRef<Path>>(path: P) -> std::io::Result<ListDir> {
    let iter = std::fs::read_dir(path)?;

    Ok(ListDir {
        iter: iter,
    })
}

pub struct WalkDir {
    // TODO: Add support for stopping at device boundaries.
    root: Option<Entry>,
    pending: Vec<ListDir>,
}

impl WalkDir {

    fn push(&mut self, entry: &Entry) {
        match list_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(iter);
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<Entry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for WalkDir {

    type Item = Entry;

    fn next(&mut self) -> Option<Entry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl std::iter::Iterator for ListDir {

    type Item = Entry;

    fn next(&mut self) -> Option<Entry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(Entry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    #[test]
    fn test_list_dir_non_existing() {
        let tempdir = tempfile::tempdir().unwrap();

        let iter = list_dir(tempdir.path().join("foo"));
        assert!(iter.is_err());
    }

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = list_dir(&tempdir).unwrap();

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_list_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();
        File::create(tempdir.path().join("што й па мору")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path().join("zażółć gęślą jaźń"));
        assert_eq!(results[1].path, tempdir.path().join("што й па мору"));
    }

    #[test]
    fn test_walk_dir_non_existing() {
        let tempdir = tempfile::tempdir().unwrap();

        let iter = walk_dir(tempdir.path().join("foo"));
        assert!(iter.is_err());
    }

    #[test]
    fn test_walk_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk_dir(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_dir_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_dir_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_dir_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_dir_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<|endoftext|>